Hi, everyone. I want a rank layer and then use backpropagation transfer its gradient to its original position, and I have written a custom function. When I test it, the forward propagation works well, and I can get expected results. However, I can not get the result of gradient. I’m really curious about where I have ignored.
This is my custom function:
class RankingFunc(Function):
@staticmethod
def forward(ctx, x, ranking_size):
p, c, h, w = x.size()
y = torch.empty(x.size())
ind = torch.arange(0, h*w).view(h, w).unsqueeze(0).repeat(p, c, 1, 1)
ctx.ranking_size = ranking_size
h_num = h // ranking_size
w_num = w // ranking_size
for i in range(h_num):
for j in range(w_num):
temp1 = x[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size].clone()
temp2 = ind[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size].clone()
value, index = torch.sort(temp1.view(p, c, ranking_size ** 2))
temp1 = value.view(p, c, ranking_size, ranking_size)
y[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size] = temp1
temp2 = temp2.view(p, c, ranking_size ** 2).view(-1)
index = index.view(-1)
temp2 = temp2[index].view(p, c, ranking_size, ranking_size)
ind[:, :, i * ranking_size : (i + 1) * ranking_size, j * ranking_size : (j + 1) * ranking_size] = temp2
ctx.save_for_backward(ind)
return y
@staticmethod
def backward(ctx, grad_y):
index = ctx.saved_variables[0]
p, c, h, w = grad_y.size()
grad_y = grad_y.view(p, c, -1).contiguous()
_, index = index.view(p, c, -1).sort()
grad_y = grad_y.view(-1)[index.view(-1)]
grad_y = grad_y.view(p, c, h, w)
return grad_y, None
And this is my test code:
import torch
import RankingFunc
x = torch.rand((2, 2, 4, 4), requires_grad = True) * 100
y = RankingFunc.apply(x, 3)
z = y / 2
z_grad = torch.rand(2, 2, 4, 4) * 100
z.backward(z_grad)
There is just no output when I print(x.grad). Does anybody have any ideas? Thanks in advance!