先上结论是可导的,但是可导的也只有被筛选出来 那一部分
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.topk(z, 2, dim=-1, largest=True, sorted=True)
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([[[[ 4., 3.],
[ 9., 8.],
[14., 13.],
[19., 18.],
[24., 23.]],
[[29., 28.],
[34., 33.],
[39., 38.],
[44., 43.],
[49., 48.]]]], grad_fn=<TopkBackward>)
zgrad tensor([[[[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.]],
[[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.],
[ 0., 0., 0., 1038., 1038.]]]])