Hi, I was wondering how can I accumulate gradient during gradient descent in pytorch (i.e. iter_size in caffe prototxt), since a single GPU can’t hold very large models now. I know here already talked about this, but I just want to confirm my code is correct. Thank you very much.
I attach my code snippets as below:
optimizer.zero_grad()
loss_mini_batch = 0

for i, (input, target) in enumerate(train_loader):

    input = input.float().cuda(async=True)
    target = target.cuda(async=True)
    input_var = torch.autograd.Variable(input)
    target_var = torch.autograd.Variable(target)

    output = model(input_var)
    loss = criterion(output, target_var)
    loss.backward()
    loss_mini_batch += loss.data[0]

    if (i+1) % args.iter_size == 0:
        # Do a SGD step once every iter_size iterations
        optimizer.step()
        optimizer.zero_grad()
        print("Iter: %02d, Loss: %4.4f" % (i, loss_mini_batch/args.iter_size))
        loss_mini_batch = 0

Thanks for any help.