def _valid_s(self, S: ModelTrainer, inputs, target):
target_ohe = to_onehot(torch.squeeze(target, dim=1).long(), num_classes=4)
target = torch.squeeze(target, dim=1).long()
seg_pred = torch.nn.functional.softmax(S.forward(inputs), dim=1)
loss_S = S.compute_loss("DiceLoss", seg_pred, target_ohe)
S.update_valid_loss("DiceLoss", loss_S.mean())
metrics = S.compute_metrics(seg_pred, target)
metrics["Dice"] = metrics["Dice"].mean()
metrics["IoU"] = metrics["IoU"].mean()
S.update_valid_metrics(metrics)
return seg_pred, loss_S
def _train_s(self, S: ModelTrainer, inputs, target, backward=True):
S.zero_grad()
target_ohe = to_onehot(torch.squeeze(target, dim=1).long(), num_classes=4)
target = torch.squeeze(target, dim=1).long()
seg_pred = torch.nn.functional.softmax(S.forward(inputs), dim=1)
loss_S = S.compute_loss("DiceLoss", seg_pred, target_ohe)
S.update_train_loss("DiceLoss", loss_S.mean())
metrics = S.compute_metrics(seg_pred, target)
metrics["Dice"] = metrics["Dice"].mean()
metrics["IoU"] = metrics["IoU"].mean()
S.update_train_metrics(metrics)
if backward:
loss_S.mean().backward()
S.step()
return seg_pred, loss_S
