How to implement gradient accumulation for GAN?

AbnerC: <p>I want to take the advantage of gradient accumulation to train a GAN with a larger batch size. I understand that for a normal network we just do something like</p>
<pre><code class="lang-auto">output = net(input)
loss = criterion(output, target_var)
loss = loss / accumulate_steps
loss.backward()
if iterations % accumulate_steps == 0:
    optimizer.step()
    optimizer.zero_grad()
</code></pre>
<p>But how can I implement this for a GAN? Cause in training a GAN we need to iteratively update G and D. When calculate the gradient of G, the wrong gradient will be accumulated to D. So we normally clear D’s gradient in each iteration which conflicts with using gradient accumulation strategy.</p>
<p>Here is a code of GAN without accumulate gradient:</p>
<pre><code class="lang-auto">#-----------
# Update G
#-----------
optimizer_G.zero_grad()
gen_imgs = generator(input_noise)
g_loss = adversarial_loss(discriminator(gen_imgs), label_real) 
g_loss.backward() 
optimizer_G.step() 

#----------
# Update D
#----------
optimizer_D.zero_grad()   # This step would clear the wrong gradient
real_loss = adversarial_loss(discriminator(real_imgs), label_real) 
fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), label_fake)
d_loss = (real_loss + fake_loss) / 2
d_loss.backward()

</code></pre>

ptrblck: <p>You could try to set the <code>.requires_grad</code> of the parameters of the model, which should not get gradients, to <code>False</code>. This should still calculate the gradients for e.g. the generator, while no new gradients should be accumulated in the discriminator.<br>
After this step, you would have to set these attribute to <code>True</code> again for the parameters.</p>

Yujian_Liu: <p><a class="mention" href="/u/abnerc">@AbnerC</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for your question and solution.<br>
<a class="mention" href="/u/abnerc">@AbnerC</a> I wonder if you have tried the mentioned solution. If I understand correctly, this means we need to set the parameters in the discriminator to False when training the generator, and set them back to True when training the discriminator, and we need to do this in every iteration.</p>

jondapper: <p>An easier solution for gradient accumulation:</p>
<pre><code class="lang-auto">#-----------
# Update G
#-----------
gen_imgs = generator(input_noise)
g_loss = adversarial_loss(discriminator(gen_imgs), label_real) 
g_loss.backward(inputs = list(generator.parameters())) 
if iter % accum_steps == 0:
    optimizer_G.step()
    optimizer_G.zero_grad()

#----------
# Update D
#----------
real_loss = adversarial_loss(discriminator(real_imgs), label_real) 
fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), label_fake)
d_loss = (real_loss + fake_loss) / 2
d_loss.backward()
if iter % accum_steps == 0:
    optimizer_D.step()
    optimizer_D.zero_grad()   
</code></pre>

pabloruizp: <p><a class="mention" href="/u/jondapper">@jondapper</a> your approach is not correct in this particular case. As to calculate the loss of the generator you need to use the discriminator, gradients will be computed for both. Each of the optimizers will update only the parameters of each module, but the gradients used will be updated in both backward passes, resulting in incorrect gradient accumulation.</p>
<p>You can check by printing the gradients of one of the modules before and after the backward pass of each of the losses. You will observe how the gradients of the discriminator will be modified when performing the backward pass of the generator loss.</p>
<p>The solution proposed by <a class="mention" href="/u/ptrblck">@ptrblck</a> avoids this problem.</p>