Variable可以把输出的Tensor变成一个输入变量,这样梯度就不会回传了。detach()也是可以的如果都不加那么就得retain_graph=True了,否则报错
import torch
import torch.nn as nn
from torch.autograd import Variable
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)
a = self.k(z)
return a
# return self.k(a)
c = 2
h = 5
w = 5
z1 = torch.arange(0, c * h * w).float().view(1, c, h, w)
z2 = torch.arange(0, c * h * w).float().view(1, c, h, w)+5
k1 = g()
k2 = g()
r1 = k1(z1)
r2 = k2(z2)
# print(r1,Variable(r1))
# loss1 = nn.MSELoss()(r1,Variable(r2))
# loss2 = nn.MSELoss()(Variable(r1),r2)
# loss1 = nn.MSELoss()(r1,r2.detach())
# loss2 = nn.MSELoss()(r1.detach(),r2)
loss1 = nn.MSELoss()(r1,r2)
loss2 = nn.MSELoss()(r1,r2)
loss1.backward()
loss2.backward()
# print("zgrad", z.grad)
输出
tensor([[[[ -8.4214, -8.9186, -9.4158, -9.9130, -10.4102],
[-10.9074, -11.4046, -11.9018, -12.3990, -12.8961],
[-13.3933, -13.8905, -14.3877, -14.8849, -15.3821],
[-15.8793, -16.3765, -16.8737, -17.3709, -17.8681],
[-18.3653, -18.8625, -19.3596, -19.8568, -20.3540]]]],
grad_fn=<MkldnnConvolutionBackward>) tensor([[[[ -8.4214, -8.9186, -9.4158, -9.9130, -10.4102],
[-10.9074, -11.4046, -11.9018, -12.3990, -12.8961],
[-13.3933, -13.8905, -14.3877, -14.8849, -15.3821],
[-15.8793, -16.3765, -16.8737, -17.3709, -17.8681],
[-18.3653, -18.8625, -19.3596, -19.8568, -20.3540]]]])