I am working on a regression problem involving a very sparse dataset (<1% non-zero examples). One of the strategies I would like to try is to accumulate the gradient for each batch over an entire epoch, and only then update the model (since a batch has many zero examples and the gradients quickly tend to zero). I am thinking of a pseudocode similar to the following:
for e in epochs:
grad = 0
for batch in dataloader:
grad += model.backward(batch)
grad /= len(dataloader)
model.update(grad)
My current code that updates the gradient after every batch is:
for e in epochs:
for batch in dataloader:
# Forward pass
optimizer.zero_grad()
x: torch.tensor = batch[0].float().to(device)
y: torch.tensor = batch[1].float().to(device)
y_pred = model_reg(x).squeeze()
loss = criterion(y_pred, y)
# Backward pass
loss.backward()
optimizer.step()
Is there a way I could do what I am thinking of in Pytorch? I know autograd is a bit special, so I wanted to seek your wisdom in these matters. My current idea was to have the optimizer.zero_grad() and optimizer.step() being done outside of the batch loops (so done only once per epoch). However, I am not sure how I could divide the accumulated gradient at the end of an epoch by the number of batches.
I have been looking through the documentation and past topics and couldn’t find an answer. Any advice or help are greatly appreciated!
Thanks,
Alex