Is grid.grad=0 when we use mode="nearest" in grid_sample?

[BailiangJ]

Hi,

I saw that in the 3D registration tutorial,

https://github.com/Project-MONAI/tutorials/blob/master/3d_registration/paired_lung_ct.ipynb

warp_layer = Warp().to(device)
pred_image = warp_layer(moving_image, ddf)
pred_label = warp_layer(moving_label, ddf)

the label is using the same warper/resampler as the image, which by default uses mode='bilinear' in torch.nn.functional.grid_sample.

Why shouldn't we use mode='nearest' here? Is it because if we use mode='nearest', the gradient wrt grid would be zero?

https://discuss.pytorch.org/t/does-the-nn-funtional-grid-sample-backwards-grad-to-grid/31736

Thanks in advance!

[wyli]

yes I think this is task dependent, linear works well in this case, could you help confirm @kate-sann5100 @YipengHu

[BailiangJ]

Yes, here is the experiment I did:

num_iteration = 5
fixed_label = 
moving_label = 
#initial dice score > 0.6

theta = torch.tensor([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], dtype=torch.float).to(fixed_label)
theta = Variable(theta, requires_grad=True)

optimizer = torch.optim.Adam([theta], lr=lr)

for iter in range(num_iteration):
        optimizer.zero_grad()
        flow = F.affine_grid(theta, moving_label.size(), align_corners=True)
        warp_moving = F.grid_sample(moving_label, grid=flow, mode="nearest", align_corners=True)

        warp_moving.retain_grad()
        flow.retain_grad()

        loss = torch.sum( (fixed_label - warp_moving)**2 ) #I use SSE loss here
        print(f"iter {iter}, loss: {loss}")
        loss.backward()

        print(torch.unique(warp_moving.grad))
        print(torch.unique(flow.grad))

        optimizer.step()

and the result I got is:

iter 0, loss: 22536.0
tensor([-2.,  0.,  2.])
tensor([0.])
iter 1, loss: 22536.0
tensor([-2.,  0.,  2.])
tensor([0.])
iter 2, loss: 22536.0
tensor([-2.,  0.,  2.])
tensor([0.])
iter 3, loss: 22536.0
tensor([-2.,  0.,  2.])
tensor([0.])
iter 4, loss: 22536.0
tensor([-2.,  0.,  2.])
tensor([0.])

the code was not doing the affine registration as I wanted (loss is not decreasing),
but when I change the mode="bilinear" , the gradient w.r.t flow is not all zero,
and the dice score could get to 0.97 (initial is around 0.66) after 200 iterations

[YipengHu]

it looks like it settles the argument - although i'm a bit surprised PyTorch didn't come up with some clever approximation for non-zero gradient - we should not only recommend but perhaps enforce the bilinear/trilinear mode. Thanks @BailiangJ

@kate-sann5100 do you think you can reproduce this test and add a note/change to the PyTorch implementation accordingly?
@wyli does the MONAI cuda resampler has non-zero gradient with nearest option?

[kate-sann5100]

@BailiangJ Thank you for the useful information.

@YipengHu I will double-check the MONAI resampler gradient and amend the PyTorch implementation next week.

[YipengHu]

also, try initialise the matrix with some small random noise to see if it still holds.

[wyli]

in my understanding, it's zeros in the MONAI implementation as well. I'm not familiar with this, but it might be interesting to see some differentiable approximations? google returned this https://math.stackexchange.com/questions/842647/polynomial-equal-to-the-ceiling-of-x