R8206/losses.py
b77bb05417bbmaster

losses.py
View Options

			import torch
			from torch import nn
			from .utils import torch_version_major, torch_version_minor

			__all__ = ["CrossEntropyLoss"]

			class CrossEntropyLoss(nn.Module):

			def __init__(self, class_weights, ignore_index=255):
			super().__init__()

			if torch_version_major==0 and torch_version_minor<=3:
			self.loss = nn.CrossEntropyLoss(weight=torch.FloatTensor(class_weights),
			ignore_index=ignore_index,
			reduce=False)
			else:
			self.loss = nn.CrossEntropyLoss(weight=torch.FloatTensor(class_weights),
			ignore_index=ignore_index,
			reduction='none')

			def forward(self, pred, target, weights=None):

			loss = self.loss(pred, target)
			if weights is not None:
			loss *= weights

			return loss.mean()

			class MSELoss(nn.Module):

			def __init__(self, ignore_index=255):
			super().__init__()
			self.ignore_index = ignore_index

			def forward(self, pred, target, weights=None):

			loss = (pred-target).pow(2)
			if weights is not None:
			loss *= weights

			if self.ignore_index is not None:
			loss = loss[target!=self.ignore_index]

			return loss.mean()