| 64 | |
| 65 | |
| 66 | def check_gradient_numerical(channels=4, grad_value=True, grad_sampling_loc=True, grad_attn_weight=True): |
| 67 | |
| 68 | value = torch.rand(N, S, M, channels).cuda() * 0.01 |
| 69 | sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda() |
| 70 | attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5 |
| 71 | attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True) |
| 72 | im2col_step = 2 |
| 73 | func = MSDeformAttnFunction.apply |
| 74 | |
| 75 | value.requires_grad = grad_value |
| 76 | sampling_locations.requires_grad = grad_sampling_loc |
| 77 | attention_weights.requires_grad = grad_attn_weight |
| 78 | |
| 79 | gradok = gradcheck(func, (value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step)) |
| 80 | |
| 81 | print(f'* {gradok} check_gradient_numerical(D={channels})') |
| 82 | |
| 83 | |
| 84 | if __name__ == '__main__': |