I work on cuda8.0 and cudnn.
The version of pytorch is 0.4.0.
I find it is very slow when I apply gradient penalty (GP) for training cifar10 with resnet18.
I test the average running time of each step with and without GP:
without : 0.065s
with GP: 0.330s
Here is my code:
import torch
from torch import nn, autograd
from torch.autograd import Variable
from models import PreActResNet18
import time

net = PreActResNet18()
net.cuda()
opt = torch.optim.SGD(net.parameters(), 0.01, momentum=0.9, weight_decay=1e-4)
batchsize = 128

GP = True
start = time.time()

for i in range(100):
    x = torch.rand((batchsize, 3, 32, 32)).cuda()
    y = torch.randint(0, 10, (batchsize,)).cuda().long()
    x, y = Variable(x, requires_grad=True), Variable(y)
    opt.zero_grad()
    preds = net(x)
    loss_c = nn.CrossEntropyLoss()(preds, y)

    if GP:
        grad = autograd.grad(loss_c, x, create_graph=True, retain_graph=True,
                             only_inputs=True)[0]
        loss_g = (grad ** 2).mean() * batchsize
        (loss_c + loss_g).backward()
    else:
        loss_c.backward()
    opt.step()

end = time.time()
print("{:.3f}s for each step".format((end-start)/100))

Is there better implementation?