@mcarilli thanks so much. I’m more of a musician than programmer so sometimes basic things need clarifying.
I’m keeping an eye on GPU memory consumption with watch nvidia-smi, and there’s barely any difference with and without calling autocast on the model’s forward pass / scaling the loss computation etc.
I’ve tried to implement autocasting as follows, currently I’m hitting memory limits before even reaching loss calculation in the training loop:
@autocast()
def forward(self, z):
out = self.l1(z)
out = out.view(out.shape[0], 128, self.init_size, self.init_size)
img = self.conv_blocks(out)
return img
g_loss = adversarial_loss(discriminator(gen_imgs), valid)
with autocast:
scaler.scale(g_loss).backward()
scaler.step(optimizer_G)
scaler.update()
It feels as though I need to recast my inputs at the beginning of the training loop to FP16 (and possibly at the transforms/dataloader stage too?), is that right?
Mark