About the loss calculation
Opened this issue · 2 comments
LordLiang commented
Maybe I think you could set
self.mse_loss = nn.MSELoss()
self.bce_loss = nn.BCELoss()
to
self.mse_loss = nn.MSELoss(size_average=False)
self.bce_loss = nn.BCELoss(size_average=False)
and then calculate losses like this way:
n_mask = torch.sum(mask)
n_noobj_mask = torch.sum(noobj_mask)
loss_x = self.bce_loss(x * mask, tx * mask) / n_mask
loss_y = self.bce_loss(y * mask, ty * mask) / n_mask
loss_w = self.mse_loss(w * mask, tw * mask) / n_mask
loss_h = self.mse_loss(h * mask, th * mask) / n_mask
loss_conf = self.bce_loss(conf * mask, 1.0 * mask) / n_mask + \
0.5 * self.bce_loss(conf * noobj_mask, 0.0 * noobj_mask) / n_noobj_mask
loss_cls = self.bce_loss(pred_cls[mask == 1], tcls[mask == 1]) / n_mask
GNAYUOHZ commented
Hi, do you test it? can you tell me the result? Thank you
chrisway613 commented
Maybe I think you could set
self.mse_loss = nn.MSELoss() self.bce_loss = nn.BCELoss()to
self.mse_loss = nn.MSELoss(size_average=False) self.bce_loss = nn.BCELoss(size_average=False)and then calculate losses like this way:
n_mask = torch.sum(mask) n_noobj_mask = torch.sum(noobj_mask) loss_x = self.bce_loss(x * mask, tx * mask) / n_mask loss_y = self.bce_loss(y * mask, ty * mask) / n_mask loss_w = self.mse_loss(w * mask, tw * mask) / n_mask loss_h = self.mse_loss(h * mask, th * mask) / n_mask loss_conf = self.bce_loss(conf * mask, 1.0 * mask) / n_mask + \ 0.5 * self.bce_loss(conf * noobj_mask, 0.0 * noobj_mask) / n_noobj_mask loss_cls = self.bce_loss(pred_cls[mask == 1], tcls[mask == 1]) / n_mask
I think maybe also add the condition when n_mask is zero:
if n_mask:
....
else:
loss_x = torch.zeros(1).cuda()
loss_y = torch.zeros(1).cuda()
loss_w = torch.zeros(1).cuda()
loss_h = torch.zeros(1).cuda()
loss_conf = self.lambda_conf * .5 * self.bce_loss(pred_conf * noobj_mask, noobj_mask * 0.) / n_noobj
loss_cls = torch.zeros(1).cuda()