def __init__(self, model, args, train_dataset, eval_dataset,
compute_metrics, loss_type, loss_gamma):
Trainer.__init__(self,
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
compute_metrics=compute_metrics)
if loss_type == 'DiceLoss':
self.loss_fct = DiceLoss()
elif loss_type == 'FocalLoss':
self.loss_fct = FocalLoss(gamma=loss_gamma)
elif loss_type == 'LabelSmoothingCrossEntropy':
self.loss_fct = LabelSmoothingCrossEntropy()
elif loss_type == 'CrossEntropyLoss':
self.loss_fct = CrossEntropyLoss()
elif loss_type == 'CourageLoss':
self.loss_fct = CourageLoss(gamma=loss_gamma)
else:
raise ValueError("Doesn't support such loss type")
