先上结论,可导的,会将梯度回传到排序前的
import torch
import torch
import torch.nn as nn
class g(nn.Module):
def __init__(self):
super(g, self).__init__()
#self.k = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=1, padding=0, bias=False)
def forward(self, z):
a,b =torch.sort(z, dim=-1)
return a
#return self.k(a)
c = 2
h = 5
w = 5
z = torch.arange(0, c * h * w).float().view(1, c, h, w)
z.requires_grad = True
k=g()
r=k(z)
print(r)
r=r.sum()
loss=(r-1)*(r-1)
loss.backward()
print("zgrad", z.grad)
tensor([[[[ 0., 1., 2., 3., 4.],
[ 5., 6., 7., 8., 9.],
[10., 11., 12., 13., 14.],
[15., 16., 17., 18., 19.],
[20., 21., 22., 23., 24.]],
[[25., 26., 27., 28., 29.],
[30., 31., 32., 33., 34.],
[35., 36., 37., 38., 39.],
[40., 41., 42., 43., 44.],
[45., 46., 47., 48., 49.]]]], grad_fn=<SortBackward>)
zgrad tensor([[[[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.]],
[[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.],
[2448., 2448., 2448., 2448., 2448.]]]])