def bce_dice_with_ssim(mask_tuple, depth_tuple):
mask_loss = (dice_loss(mask_tuple[0], mask_tuple[1]) +
LocalBCELoss(mask_tuple[0], mask_tuple[1])) / 2
depth_loss = ssim(depth_tuple[0], depth_tuple[1], 11, reduction='mean')
#torch.sqrt(nn.MSELoss()(depth_tuple[0], depth_tuple[1]))
#ssim(depth_tuple[0], depth_tuple[1], 11)
return mask_loss, depth_loss
def bce_with_rmse_ssim(mask_tuple, depth_tuple):
mask_loss = LocalBCELoss(mask_tuple[0], mask_tuple[1])
depth_loss = (1 - ssim(
depth_tuple[0], depth_tuple[1], 11, reduction='mean')) + torch.sqrt(
nn.MSELoss()(depth_tuple[0], depth_tuple[1]))
#torch.sqrt(nn.MSELoss()(depth_tuple[0], depth_tuple[1]))
#ssim(depth_tuple[0], depth_tuple[1], 11)
return mask_loss, depth_loss
def bce_dice_with_rmse_ssim(mask_tuple, depth_tuple):
mask_loss = (dice_loss(mask_tuple[0], mask_tuple[1]) +
LocalBCELoss(mask_tuple[0], mask_tuple[1])) / 2
depth_loss = (torch.sqrt(nn.MSELoss()(depth_tuple[0], depth_tuple[1])) +
ssim(depth_tuple[0], depth_tuple[1], 11)) / 2
#torch.sqrt(nn.MSELoss()(depth_tuple[0], depth_tuple[1]))
#ssim(depth_tuple[0], depth_tuple[1], 11)
return mask_loss, depth_loss
print("[%d,%d] Loss:%.4f, PSNR: %.4f, SSIM: %.4f" %
(epoch + 1, i + 1, loss.item(), psnr_val, ssim_val))
#3. Log the scalar values
writer.add_scalar('loss', loss.item(), step)
step += 1
''' validation '''
psnr_val = 0.
ssim_val = 0.
with torch.no_grad():
net.eval()
for i, (lr, guide, gt) in enumerate(loader['validation']):
lr, guide, gt = lr.cuda(), guide.cuda(), gt.cuda()
lr = F.interpolate(lr, size=guide.shape[-2:])
imgf = torch.clamp(net(lr, guide), 0., 1.)
psnr_val += psnr_loss(imgf, gt, 1.)
ssim_val += ssim(imgf, gt, 5, 'mean', 1.)
psnr_val = float(psnr_val / loader['validation'].__len__())
ssim_val = 1 - 2 * float(ssim_val / loader['validation'].__len__())
writer.add_scalar('PSNR on validation data', psnr_val, epoch)
writer.add_scalar('SSIM on validation data', ssim_val, epoch)
''' decay the learning rate '''
# scheduler.step()
''' save model '''
if best_psnr_val < psnr_val:
best_psnr_val = psnr_val
torch.save(net.state_dict(), os.path.join(save_path, 'best_net.pth'))
torch.save(
{
'net': net.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch
def LocalSSIMLoss(input, target, reduction='mean'):
return ssim(input, target, 11, reduction=reduction)
def joint_ssim_loss(mask_tuple, depth_tuple):
mask_ssim_loss = ssim(mask_tuple[0], mask_tuple[1], 11)
depth_ssim_loss = ssim(depth_tuple[0], depth_tuple[1], 11)
return (mask_ssim_loss + depth_ssim_loss) / 2
