lyakaap / VAT-pytorch

Virtual Adversarial Training (VAT) implementation for PyTorch
296 stars 43 forks source link

VAT-pytorch

Virtual Adversarial Training (VAT) implementation for Pytorch

Usage

for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    optimizer.zero_grad()

    vat_loss = VATLoss(xi=10.0, eps=1.0, ip=1)
    cross_entropy = nn.CrossEntropyLoss()

    # LDS should be calculated before the forward for cross entropy
    lds = vat_loss(model, data)
    output = model(data)
    loss = cross_entropy(output, target) + args.alpha * lds
    loss.backward()
    optimizer.step()