先说结论不需要,其他的自定义函数也行的
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):
return self.k(z)
def encode(self, z):
return self.k(z)
c = 2
h = 5
w = 5
z = torch.rand( (1,c , h , w)).float().view(1, c, h, w)*100
z.requires_grad = True
k = g()
r = k.encode(z)
r = r.sum()
loss = (r - 1) * (r - 1)
for name,v in k.named_parameters():
print(name,v.grad)
print("*********************")
loss.backward()
for name,v in k.named_parameters():
print(name,v.grad)
结果
k.weight None
***************************
k.weight tensor([[[[426972.1875]],
[[383235.8438]]]])