I'm playing with different reduction methods provided in built-in loss functions. In particular, I would like to compare the following.
The averaged gradient by performing backward pass for each loss value calculated with
reduction="none"
The gradient averaged by dividing the batch size with
reduction="sum"
The average gradient yielded by
reduction="mean"
- The average gradient calculated by
reduction="mean"
, with the data points fed into the model one at a time.
My code for producing the experiment is as follows:
def estimate_gradient(model, optimizer, batch):
criterion_no_reduction = nn.CrossEntropyLoss(reduction="none").cuda()
criterion_sum = nn.CrossEntropyLoss(reduction="sum").cuda()
criterion_avg = nn.CrossEntropyLoss().cuda()
input, target = batch
input, target = input.cuda(), target.cuda()
output = model(input)
n = len(output)
loss_no_reudction = criterion_no_reduction(output, target)
grad_list_no_reduction = []
for i in range(n):
optimizer.zero_grad()
loss_no_reudction[i].backward(retain_graph=True)
for j, param in enumerate(model.parameters()):
if param.requires_grad:
grad = param.grad.view(-1, 1)
if i == 0:
grad_list_no_reduction.append(grad)
else:
grad_list_no_reduction[j] = torch.cat((grad_list_no_reduction[j], grad), dim=1)
grad_out_no_reduction = torch.cat(grad_list_no_reduction, dim=0)
grad_out_no_reduction = (torch.sum(grad_out_no_reduction, dim=1) / n).cpu().detach().numpy().flatten()
loss_sum = criterion_sum(output, target)
optimizer.zero_grad()
loss_sum.backward(retain_graph=True)
for j, param in enumerate(model.parameters()):
if param.requires_grad:
if j == 0:
grad_list_sum = param.grad.view(-1)
else:
grad_list_sum = torch.cat((grad_list_sum, param.grad.view(-1)))
grad_out_sum = (grad_list_sum / n).cpu().detach().numpy().flatten()
loss_avg = criterion_avg(output, target)
optimizer.zero_grad()
loss_avg.backward(retain_graph=True)
for j, param in enumerate(model.parameters()):
if param.requires_grad:
if j == 0:
grad_list_avg = param.grad.view(-1)
else:
grad_list_avg = torch.cat((grad_list_avg, param.grad.view(-1)))
grad_out_avg = grad_list_avg.cpu().detach().numpy().flatten()
target = target.view(-1, 1)
grad_list_one_by_one = []
for i in range(n):
optimizer.zero_grad()
curr_output = output[i].view(1, -1)
loss = criterion_avg(curr_output, target[i])
loss.backward(retain_graph=True)
for j, param in enumerate(model.parameters()):
if param.requires_grad:
grad = param.grad.view(-1, 1)
if i == 0:
grad_list_one_by_one.append(grad)
else:
grad_list_one_by_one[j] = torch.cat((grad_list_one_by_one[j], grad), dim=1)
grad_out_one_by_one = torch.cat(grad_list_one_by_one, dim=0)
grad_out_one_by_one = (torch.sum(grad_out_one_by_one, dim=1) / n).cpu().detach().numpy().flatten()
assert grad_out_no_reduction.shape == grad_out_sum.shape == grad_out_avg.shape == grad_out_one_by_one.shape
print("Maximum discrepancy between reduction = none and sum: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_sum))))
print("Maximum discrepancy between reduction = none and avg: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_avg))))
print("Maximum discrepancy between reduction = none and one-by-one: {}".format(np.max(np.abs(grad_out_no_reduction - grad_out_one_by_one))))
print("Maximum discrepancy between reduction = sum and avg: {}".format(np.max(np.abs(grad_out_sum - grad_out_avg))))
print("Maximum discrepancy between reduction = sum and one-by-one: {}".format(np.max(np.abs(grad_out_sum - grad_out_one_by_one))))
print("Maximum discrepancy between reduction = avg and one-by-one: {}".format(np.max(np.abs(grad_out_avg- grad_out_one_by_one))))
The results are as follows:
Maximum discrepancy between reduction = none and sum: 0.0316
Maximum discrepancy between reduction = none and avg: 0.0316
Maximum discrepancy between reduction = none and one-by-one: 0.0
Maximum discrepancy between reduction = sum and avg: 0.0
Maximum discrepancy between reduction = sum and one-by-one: 0.0316
Maximum discrepancy between reduction = avg and one-by-one: 0.0316
That is, the result produced by reduction=none
and one-by-one backward pass appear to be identical, while reduciton=sum
and reduction=mean
yields different results from the previous pair. It would be really helpful to explain the discrepancy (maybe due to retain_graph=True
?) and thanks in advance for any help!