File size: 2,891 Bytes
a639402 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import torch.nn as nn
import segmentation_models_pytorch as smp
import torch
import torch.nn as nn
import torch.nn.functional as F
'''
Segmentation Model Losses
Source: https://smp.readthedocs.io/en/latest/losses.html
'''
def set_classification_loss(weight, loss_function='CE'):
if loss_function == 'MM':
criterion = nn.MultiMarginLoss()
elif loss_function == 'CE':
criterion = nn.CrossEntropyLoss(weight=weight)
return criterion
def set_segmentation_loss(loss_function='Dice',
mode='binary',
num_classes=None,
log_loss=False,
from_logits=True
):
if loss_function == 'Dice':
criterion = smp.losses.DiceLoss(mode=mode,
classes=None,
log_loss=log_loss,
from_logits=from_logits)
elif loss_function == 'Jaccard':
criterion = smp.losses.DiceLoss(mode=mode,
classes=num_classes,
log_loss=log_loss,
from_logits=from_logits,
smooth=0.0,
eps=1e-07)
return criterion
class ContrastiveLoss(nn.Module):
def __init__(self, margin=1.0):
super(ContrastiveLoss, self).__init__()
self.margin = margin
def forward(self, input_embed, pos_embed, neg_embed):
distance_pos = F.pairwise_distance(input_embed, pos_embed, p=2)
distance_neg = F.pairwise_distance(input_embed, neg_embed, p=2)
loss = torch.mean(torch.clamp(self.margin + distance_pos - distance_neg, min=0.0))
return loss
class WeightedContrastiveLoss(nn.Module):
def __init__(self, margin=1.0):
super(WeightedContrastiveLoss, self).__init__()
self.margin = margin
def forward(self, input_embed, pos_embed, neg_embed, weight=None):
distance_pos = F.pairwise_distance(input_embed, pos_embed, p=2)
distance_neg = F.pairwise_distance(input_embed, neg_embed, p=2)
loss = torch.clamp(self.margin + distance_pos - distance_neg, min=0.0)
if weight is not None:
loss = (loss * weight / weight.sum()).sum()
loss = torch.mean(loss)
return loss
class WeightedCrossEntropy(nn.Module):
def __init__(self):
super(WeightedCrossEntropy, self).__init__()
def forward(self, weight=None):
loss = nn.CrossEntropyLoss(reduction='none')
if weight is not None:
loss =(loss * weight / weight.sum()).sum()
loss = loss.mean()
#loss = torch.mean(loss)
return loss |