MLNet-Pytorch
MLNet-Pytorch copied to clipboard
ModMSELoss is incorrect
Thanks for your great work! But there seems to be an incorrect implementation for the ModMSELoss.
- The division operation
(output / output_max)
could causeNaN
values if all entries ofoutput
are zeros. - The square operation should take the whole fraction term as input.
Therefore, I changed the forward function of ModMSELoss as follows:
def forward(self, output , label , prior):
prior_size = prior.shape
output_max = torch.max(torch.max(output,2)[0],2)[0].unsqueeze(2).unsqueeze(2).expand(output.shape[0],output.shape[1],self.shape_r_gt,self.shape_c_gt)
reg = ( 1.0/(prior_size[0]*prior_size[1]) ) * ( 1 - prior)**2 # (1, 1, 6, 8)
loss = torch.mean( ((output / (output_max + 1e-6) - label) / (1 - label + 0.1))**2) + torch.sum(reg)
return loss