import torch from torch import nn from torch.nn.functional import _pointwise_loss rgb_weights = [0.29891 * 3, 0.58661 * 3, 0.11448 * 3] # RGB have different weights # https://github.com/nagadomi/waifu2x/blob/master/train.lua#L109 use_cuda = torch.cuda.is_available() FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor LongTensor = torch.cuda.LongTensor if use_cuda else torch.LongTensor Tensor = FloatTensor class WeightedHuberLoss(nn.SmoothL1Loss): def __init__(self, weights=rgb_weights): super(WeightedHuberLoss, self).__init__(size_average=True, reduce=True) self.weights = torch.FloatTensor(weights).view(3, 1, 1) def forward(self, input_data, target): diff = torch.abs(input_data - target) z = torch.where(diff < 1, 0.5 * torch.pow(diff, 2), (diff - 0.5)) out = z * self.weights.expand_as(diff) return out.mean() def weighted_mse_loss(input, target, weights): out = (input - target) ** 2 out = out * weights.expand_as(out) loss = out.sum(0) # or sum over whatever dimensions return loss / out.size(0) class WeightedL1Loss(nn.SmoothL1Loss): def __init__(self, weights=rgb_weights): super(WeightedHuberLoss, self).__init__(size_average=True, reduce=True) self.weights = torch.FloatTensor(weights).view(3, 1, 1) def forward(self, input_data, target): return self.l1_loss( input_data, target, size_average=self.size_average, reduce=self.reduce ) def l1_loss(self, input_data, target, size_average=True, reduce=True): return _pointwise_loss( lambda a, b: torch.abs(a - b) * self.weights.expand_as(a), torch._C._nn.l1_loss, input_data, target, size_average, reduce, )