I defined a variable(self.A) in the MyLoss for segmentation task, but its value will not be updated during the training process, and no gradient value will be generated. What should I do to update the value of Variable?
class MyLoss(nn.Module):
def __init__(self):
super().__init__()
self.A = torch.nn.Parameter(torch.ones(16, 1, 62, 126, requires_grad=True).cuda())
def my_loss(self, input, target):
x = input[:, :, 1:, :] - input[:, :, :-1, :]
y = input[:, :, :, 1:] - input[:, :, :, :-1]
delta_x = x[:, :, 1:, :-2]**2
delta_y = y[:, :, :-2, 1:]**2
delta_u = torch.abs(delta_x + delta_y)
loss = torch.mean(delta_u * self.A)
return loss
def forward(self, input, target):
input_1 = F.softmax(input, dim=1)[:,1:,:,:]
my_loss = self.my_loss(input_1, target)
print(self.A)
return my_loss
data = get_databunch(img_path='...', lbl_path='...')
learn = unet_learner(data, models.resnet34, metrics=metrics, loss_func=MyLoss(), wd=1e-2)
learn.fit_one_cycle(1, slice(3e-3), pct_start=0.8)