def test_vif_loss_zero_for_equal_tensors(x):
loss = VIFLoss()
y = x.clone()
measure = loss(x, y)
assert torch.isclose(
measure, torch.tensor(0.),
atol=1e-6), f'VIF for equal tensors must be 0, got {measure}'
def test_vif_loss_zero_for_equal_tensors(prediction: torch.Tensor):
loss = VIFLoss()
target = prediction.clone()
measure = loss(prediction, target)
assert torch.isclose(
measure, torch.tensor(0.),
atol=1e-6), f'VIF for equal tensors must be 0, got {measure}'
def test_vif_loss_reduction(x, y) -> None:
loss = VIFLoss(reduction='mean')
measure = loss(x, y)
assert measure.dim() == 0, f'VIF with `mean` reduction must return 1 number, got {len(measure)}'
loss = VIFLoss(reduction='sum')
measure = loss(x, y)
assert measure.dim() == 0, f'VIF with `mean` reduction must return 1 number, got {len(measure)}'
loss = VIFLoss(reduction='none')
measure = loss(x, y)
assert len(measure) == x.size(0), \
f'VIF with `none` reduction must have length equal to number of images, got {len(measure)}'
loss = VIFLoss(reduction='random string')
with pytest.raises(ValueError):
loss(x, y)
def test_vif_loss_reduction(prediction: torch.Tensor,
target: torch.Tensor) -> None:
loss = VIFLoss(reduction='mean')
measure = loss(prediction, target)
assert measure.dim(
) == 0, f'VIF with `mean` reduction must return 1 number, got {len(measure)}'
loss = VIFLoss(reduction='sum')
measure = loss(prediction, target)
assert measure.dim(
) == 0, f'VIF with `mean` reduction must return 1 number, got {len(measure)}'
loss = VIFLoss(reduction='none')
measure = loss(prediction, target)
assert len(measure) == prediction.size(0), \
f'VIF with `none` reduction must have length equal to number of images, got {len(measure)}'
loss = VIFLoss(reduction='random string')
with pytest.raises(KeyError):
loss(prediction, target)
def test_vif_loss_computes_grad_for_zeros_tensors() -> None:
x = torch.zeros(4, 3, 256, 256, requires_grad=True)
y = torch.zeros(4, 3, 256, 256)
loss_value = VIFLoss()(x, y)
loss_value.backward()
assert x.grad is not None, NONE_GRAD_ERR_MSG
def test_vif_loss_computes_grad(x, y, device: str) -> None:
x.requires_grad_()
loss_value = VIFLoss()(x.to(device), y.to(device))
loss_value.backward()
assert x.grad is not None, NONE_GRAD_ERR_MSG
def test_vif_loss_forward(x, y, device: str) -> None:
loss = VIFLoss()
loss(x.to(device), y.to(device))
def train_net(params):
# Initialize Parameters
params = DotDict(params)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
verbose = {}
verbose['loss_train'], verbose['loss_valid'], verbose['psnr_train'], verbose['psnr_valid'], \
verbose['ssim_train'], verbose['ssim_valid'], verbose['vif_train'], verbose['vif_valid'] = ([] for i in range(8))
log_metrics = True
ssim_module = SSIM()
msssim_module = MSSSIM()
vifLoss = VIFLoss(sigma_n_sq=0.4, data_range=1.)
msssimLoss = MultiScaleSSIMLoss(data_range=1.)
best_validation_metrics = 100
train_generator, val_generator = data_loaders(params)
loaders = {"train": train_generator, "valid": val_generator}
wnet_identifier = params.mask_URate[0:2] + "WNet_dense=" + str(int(params.dense)) + "_" + params.architecture + "_" \
+ params.lossFunction + '_lr=' + str(params.lr) + '_ep=' + str(params.num_epochs) + '_complex=' \
+ str(int(params.complex_net)) + '_' + 'edgeModel=' + str(int(params.edge_model)) \
+ '(' + str(params.num_edge_slices) + ')_date=' + (datetime.now()).strftime("%d-%m-%Y_%H-%M-%S")
if not os.path.isdir(params.model_save_path):
os.mkdir(params.model_save_path)
print("\n\nModel will be saved at:\n", params.model_save_path)
print("WNet ID: ", wnet_identifier)
wnet, optimizer, best_validation_loss, preTrainedEpochs = generate_model(
params, device)
# data = (iter(train_generator)).next()
# Adding writer for tensorboard. Also start tensorboard, which tries to access logs in the runs directory
writer = init_tensorboard(iter(train_generator), wnet, wnet_identifier,
device)
for epoch in trange(preTrainedEpochs, params.num_epochs):
for phase in ['train', 'valid']:
if phase == 'train':
wnet.train()
else:
wnet.eval()
for i, data in enumerate(loaders[phase]):
# for i in range(10000):
x, y_true, _, _, fname, slice_num = data
x, y_true = x.to(device, dtype=torch.float), y_true.to(
device, dtype=torch.float)
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'):
y_pred = wnet(x)
if params.lossFunction == 'mse':
loss = F.mse_loss(y_pred, y_true)
elif params.lossFunction == 'l1':
loss = F.l1_loss(y_pred, y_true)
elif params.lossFunction == 'ssim':
# standard SSIM
loss = 0.16 * F.l1_loss(y_pred, y_true) + 0.84 * (
1 - ssim_module(y_pred, y_true))
elif params.lossFunction == 'msssim':
# loss = 0.16 * F.l1_loss(y_pred, y_true) + 0.84 * (1 - msssim_module(y_pred, y_true))
prediction_abs = torch.sqrt(
torch.square(y_pred[:, 0::2]) +
torch.square(y_pred[:, 1::2]))
target_abs = torch.sqrt(
torch.square(y_true[:, 0::2]) +
torch.square(y_true[:, 1::2]))
prediction_abs_flat = (torch.flatten(
prediction_abs, start_dim=0,
end_dim=1)).unsqueeze(1)
target_abs_flat = (torch.flatten(
target_abs, start_dim=0, end_dim=1)).unsqueeze(1)
loss = msssimLoss(prediction_abs_flat, target_abs_flat)
elif params.lossFunction == 'vif':
prediction_abs = torch.sqrt(
torch.square(y_pred[:, 0::2]) +
torch.square(y_pred[:, 1::2]))
target_abs = torch.sqrt(
torch.square(y_true[:, 0::2]) +
torch.square(y_true[:, 1::2]))
prediction_abs_flat = (torch.flatten(
prediction_abs, start_dim=0,
end_dim=1)).unsqueeze(1)
target_abs_flat = (torch.flatten(
target_abs, start_dim=0, end_dim=1)).unsqueeze(1)
loss = vifLoss(prediction_abs_flat, target_abs_flat)
elif params.lossFunction == 'mse+vif':
prediction_abs = torch.sqrt(
torch.square(y_pred[:, 0::2]) +
torch.square(y_pred[:, 1::2])).to(device)
target_abs = torch.sqrt(
torch.square(y_true[:, 0::2]) +
torch.square(y_true[:, 1::2])).to(device)
prediction_abs_flat = (torch.flatten(
prediction_abs, start_dim=0,
end_dim=1)).unsqueeze(1)
target_abs_flat = (torch.flatten(
target_abs, start_dim=0, end_dim=1)).unsqueeze(1)
loss = 0.15 * F.mse_loss(
prediction_abs_flat,
target_abs_flat) + 0.85 * vifLoss(
prediction_abs_flat, target_abs_flat)
elif params.lossFunction == 'l1+vif':
prediction_abs = torch.sqrt(
torch.square(y_pred[:, 0::2]) +
torch.square(y_pred[:, 1::2])).to(device)
target_abs = torch.sqrt(
torch.square(y_true[:, 0::2]) +
torch.square(y_true[:, 1::2])).to(device)
prediction_abs_flat = (torch.flatten(
prediction_abs, start_dim=0,
end_dim=1)).unsqueeze(1)
target_abs_flat = (torch.flatten(
target_abs, start_dim=0, end_dim=1)).unsqueeze(1)
loss = 0.146 * F.l1_loss(
y_pred, y_true) + 0.854 * vifLoss(
prediction_abs_flat, target_abs_flat)
elif params.lossFunction == 'msssim+vif':
prediction_abs = torch.sqrt(
torch.square(y_pred[:, 0::2]) +
torch.square(y_pred[:, 1::2])).to(device)
target_abs = torch.sqrt(
torch.square(y_true[:, 0::2]) +
torch.square(y_true[:, 1::2])).to(device)
prediction_abs_flat = (torch.flatten(
prediction_abs, start_dim=0,
end_dim=1)).unsqueeze(1)
target_abs_flat = (torch.flatten(
target_abs, start_dim=0, end_dim=1)).unsqueeze(1)
loss = 0.66 * msssimLoss(
prediction_abs_flat,
target_abs_flat) + 0.33 * vifLoss(
prediction_abs_flat, target_abs_flat)
if not math.isnan(loss.item()) and loss.item(
) < 2 * best_validation_loss: # avoid nan/spike values
verbose['loss_' + phase].append(loss.item())
writer.add_scalar(
'Loss/' + phase + '_epoch_' + str(epoch),
loss.item(), i)
if log_metrics and (
(i % params.verbose_gap == 0) or
(phase == 'valid' and epoch > params.verbose_delay)):
y_true_copy = y_true.detach().cpu().numpy()
y_pred_copy = y_pred.detach().cpu().numpy()
y_true_copy = y_true_copy[:, ::
2, :, :] + 1j * y_true_copy[:,
1::
2, :, :]
y_pred_copy = y_pred_copy[:, ::
2, :, :] + 1j * y_pred_copy[:,
1::
2, :, :]
if params.architecture[-1] == 'k':
# transform kspace to image domain
y_true_copy = np.fft.ifft2(y_true_copy,
axes=(2, 3))
y_pred_copy = np.fft.ifft2(y_pred_copy,
axes=(2, 3))
# Sum of squares
sos_true = np.sqrt(
(np.abs(y_true_copy)**2).sum(axis=1))
sos_pred = np.sqrt(
(np.abs(y_pred_copy)**2).sum(axis=1))
'''
# Normalization according to: extract_challenge_metrics.ipynb
sos_true_max = sos_true.max(axis = (1,2),keepdims = True)
sos_true_org = sos_true/sos_true_max
sos_pred_org = sos_pred/sos_true_max
# Normalization by normalzing with ref with max_ref and rec with max_rec, respectively
sos_true_max = sos_true.max(axis = (1,2),keepdims = True)
sos_true_mod = sos_true/sos_true_max
sos_pred_max = sos_pred.max(axis = (1,2),keepdims = True)
sos_pred_mod = sos_pred/sos_pred_max
'''
'''
# normalization by mean and std
std = sos_pred.std(axis=(1, 2), keepdims=True)
mean = sos_pred.mean(axis=(1, 2), keepdims=True)
sos_pred_std = (sos_pred-mean) / std
std = sos_true.std(axis=(1, 2), keepdims=True)
mean = sos_pred.mean(axis=(1, 2), keepdims=True)
sos_true_std = (sos_true-mean) / std
'''
'''
ssim, psnr, vif = metrics(sos_pred_org, sos_true_org)
ssim_mod, psnr_mod, vif_mod = metrics(sos_pred_mod, sos_true_mod)
'''
sos_true_max = sos_true.max(axis=(1, 2), keepdims=True)
sos_true_org = sos_true / sos_true_max
sos_pred_org = sos_pred / sos_true_max
ssim, psnr, vif = metrics(sos_pred, sos_true)
ssim_normed, psnr_normed, vif_normed = metrics(
sos_pred_org, sos_true_org)
verbose['ssim_' + phase].append(np.mean(ssim_normed))
verbose['psnr_' + phase].append(np.mean(psnr_normed))
verbose['vif_' + phase].append(np.mean(vif_normed))
'''
print("===Normalization according to: extract_challenge_metrics.ipynb===")
print("SSIM: ", verbose['ssim_'+phase][-1])
print("PSNR: ", verbose['psnr_'+phase][-1])
print("VIF: ", verbose['vif_' +phase][-1])
print("===Normalization by normalzing with ref with max_ref and rec with max_rec, respectively===")
print("SSIM_mod: ", np.mean(ssim_mod))
print("PSNR_mod: ", np.mean(psnr_mod))
print("VIF_mod: ", np.mean(vif_mod))
print("===Normalization by dividing by the standard deviation of ref and rec, respectively===")
'''
print("Epoch: ", epoch)
print("SSIM: ", np.mean(ssim))
print("PSNR: ", np.mean(psnr))
print("VIF: ", np.mean(vif))
print("SSIM_normed: ", verbose['ssim_' + phase][-1])
print("PSNR_normed: ", verbose['psnr_' + phase][-1])
print("VIF_normed: ", verbose['vif_' + phase][-1])
'''
if True: #verbose['vif_' + phase][-1] < 0.4:
plt.figure(figsize=(9, 6), dpi=150)
gs1 = gridspec.GridSpec(3, 2)
gs1.update(wspace=0.002, hspace=0.1)
plt.subplot(gs1[0])
plt.imshow(sos_true[0], cmap="gray")
plt.axis("off")
plt.subplot(gs1[1])
plt.imshow(sos_pred[0], cmap="gray")
plt.axis("off")
plt.show()
# plt.pause(10)
# plt.close()
'''
writer.add_scalar(
'SSIM/' + phase + '_epoch_' + str(epoch),
verbose['ssim_' + phase][-1], i)
writer.add_scalar(
'PSNR/' + phase + '_epoch_' + str(epoch),
verbose['psnr_' + phase][-1], i)
writer.add_scalar(
'VIF/' + phase + '_epoch_' + str(epoch),
verbose['vif_' + phase][-1], i)
print('Loss ' + phase + ': ', loss.item())
if phase == 'train':
if loss.item() < 2 * best_validation_loss:
loss.backward()
optimizer.step()
# Calculate Averages
psnr_mean = np.mean(verbose['psnr_valid'])
ssim_mean = np.mean(verbose['ssim_valid'])
vif_mean = np.mean(verbose['vif_valid'])
validation_metrics = 0.2 * psnr_mean + 0.4 * ssim_mean + 0.4 * vif_mean
valid_avg_loss_of_current_epoch = np.mean(verbose['loss_valid'])
writer.add_scalar('AvgLoss/+train_epoch_' + str(epoch),
np.mean(verbose['loss_train']), epoch)
writer.add_scalar('AvgLoss/+valid_epoch_' + str(epoch),
np.mean(verbose['loss_valid']), epoch)
writer.add_scalar('AvgSSIM/+train_epoch_' + str(epoch),
np.mean(verbose['ssim_train']), epoch)
writer.add_scalar('AvgSSIM/+valid_epoch_' + str(epoch), ssim_mean,
epoch)
writer.add_scalar('AvgPSNR/+train_epoch_' + str(epoch),
np.mean(verbose['psnr_train']), epoch)
writer.add_scalar('AvgPSNR/+valid_epoch_' + str(epoch), psnr_mean,
epoch)
writer.add_scalar('AvgVIF/+train_epoch_' + str(epoch),
np.mean(verbose['vif_train']), epoch)
writer.add_scalar('AvgVIF/+valid_epoch_' + str(epoch), vif_mean, epoch)
verbose['loss_train'], verbose['loss_valid'], verbose['psnr_train'], verbose['psnr_valid'], \
verbose['ssim_train'], verbose['ssim_valid'], verbose['vif_train'], verbose['vif_valid'] = ([] for i in
range(8))
# Save Networks/Checkpoints
if best_validation_metrics > validation_metrics:
best_validation_metrics = validation_metrics
best_validation_loss = valid_avg_loss_of_current_epoch
save_checkpoint(
wnet, params.model_save_path, wnet_identifier, {
'epoch': epoch + 1,
'state_dict': wnet.state_dict(),
'best_validation_loss': best_validation_loss,
'optimizer': optimizer.state_dict(),
}, True)
else:
save_checkpoint(
wnet, params.model_save_path, wnet_identifier, {
'epoch': epoch + 1,
'state_dict': wnet.state_dict(),
'best_validation_loss': best_validation_loss,
'optimizer': optimizer.state_dict(),
}, False)
def test_vif_loss_computes_grad_for_zeros_tensors() -> None:
prediction = torch.zeros(4, 3, 256, 256, requires_grad=True)
target = torch.zeros(4, 3, 256, 256)
loss_value = VIFLoss()(prediction, target)
loss_value.backward()
assert prediction.grad is not None, NONE_GRAD_ERR_MSG
def test_vif_loss_computes_grad(prediction: torch.Tensor, target: torch.Tensor,
device: str) -> None:
prediction.requires_grad_()
loss_value = VIFLoss()(prediction.to(device), target.to(device))
loss_value.backward()
assert prediction.grad is not None, NONE_GRAD_ERR_MSG
def test_vif_loss_forward(prediction: torch.Tensor, target: torch.Tensor,
device: str) -> None:
loss = VIFLoss()
loss(prediction.to(device), target.to(device))
def test_vif_loss_forward_on_gpu(prediction: torch.Tensor,
target: torch.Tensor) -> None:
loss = VIFLoss()
loss(prediction.cuda(), target.cuda())
def test_vif_loss_forward(prediction: torch.Tensor,
target: torch.Tensor) -> None:
loss = VIFLoss()
loss(prediction, target)
def test_vif_loss_computes_grad_on_gpu(prediction: torch.Tensor,
target: torch.Tensor) -> None:
prediction.requires_grad_()
loss_value = VIFLoss()(prediction.cuda(), target.cuda())
loss_value.backward()
assert prediction.grad is not None, NONE_GRAD_ERR_MSG