Upload models/loss.py

models/loss.py

@@ -0,0 +1,222 @@
+from __future__ import division
+import os, glob, shutil, math, random, json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import basic
+from utils import util
+
+eps = 0.0000001
+
+class SPixelLoss:
+    def __init__(self, psize=8, mpdist=False, gpu_no=0):
+        self.mpdist = mpdist
+        self.gpu_no = gpu_no
+        self.sp_size = psize
+    
+    def __call__(self, data, epoch_no):
+        kernel_size = self.sp_size
+        #pos_weight = 0.003
+        prob = data['pred_prob']
+        labxy_feat = data['target_feat']
+        N,C,H,W = labxy_feat.shape
+        pooled_labxy = basic.poolfeat(labxy_feat, prob, kernel_size, kernel_size)
+        reconstr_feat = basic.upfeat(pooled_labxy, prob, kernel_size, kernel_size)
+        loss_map = reconstr_feat[:,:,:,:] - labxy_feat[:,:,:,:]
+        featLoss_idx = torch.norm(loss_map[:,:-2,:,:], p=2, dim=1).mean()
+        posLoss_idx = torch.norm(loss_map[:,-2:,:,:], p=2, dim=1).mean() / kernel_size
+        totalLoss_idx = 10*featLoss_idx + 0.003*posLoss_idx
+        return {'totalLoss':totalLoss_idx, 'featLoss':featLoss_idx, 'posLoss':posLoss_idx}
+
+
+class AnchorColorProbLoss:
+    def __init__(self, hint2regress=False, enhanced=False, with_grad=False, mpdist=False, gpu_no=0):
+        self.mpdist = mpdist
+        self.gpu_no = gpu_no
+        self.hint2regress = hint2regress
+        self.enhanced = enhanced
+        self.with_grad = with_grad
+        self.rebalance_gradient = basic.RebalanceLoss.apply
+        self.entropy_loss = nn.CrossEntropyLoss(ignore_index=-1)
+        if self.enhanced:
+            self.VGGLoss = VGG19Loss(gpu_no=gpu_no, is_ddp=mpdist)
+    
+    def _perceptual_loss(self, input_grays, input_colors, pred_colors):
+        input_RGBs = basic.lab2rgb(torch.cat([input_grays,input_colors], dim=1))
+        pred_RGBs = basic.lab2rgb(torch.cat([input_grays,pred_colors], dim=1))
+        ## the output of "lab2rgb" just matches the input of "VGGLoss": [0,1]
+        return self.VGGLoss(input_RGBs, pred_RGBs)
+    
+    def _laplace_gradient(self, pred_AB, target_AB):
+        N,C,H,W = pred_AB.shape
+        kernel = torch.tensor([[1, 1, 1], [1, -8, 1], [1, 1, 1]], device=pred_AB.get_device()).float()
+        kernel = kernel.view(1, 1, *kernel.size()).repeat(C,1,1,1)
+        grad_pred = F.conv2d(pred_AB, kernel, groups=C)
+        grad_trg = F.conv2d(target_AB, kernel, groups=C)
+        return l1_loss(grad_trg, grad_pred)
+        
+    def __call__(self, data, epoch_no):
+        N,C,H,W = data['target_label'].shape
+        pal_probs = self.rebalance_gradient(data['pal_prob'], data['class_weight'])
+        #ref_probs = data['ref_prob']
+        pal_probs = pal_probs.permute(0,2,3,1).contiguous().view(N*H*W, -1)
+        gt_labels = data['target_label'].permute(0,2,3,1).contiguous().view(N*H*W, -1)
+        '''        
+        igored_mask = data['empty_entries'].permute(0,2,3,1).contiguous().view(N*H*W, -1)
+        gt_labels[igored_mask] = -1
+        gt_labels = gt_probs.squeeze()
+        '''
+        palLoss_idx = self.entropy_loss(pal_probs, gt_labels.squeeze(dim=1))
+        if self.hint2regress: 
+            ref_probs = data['ref_prob']
+            refLoss_idx = 50 * l2_loss(data['spix_color'], ref_probs)
+        else:
+            ref_probs = self.rebalance_gradient(data['ref_prob'], data['class_weight'])
+            ref_probs = ref_probs.permute(0,2,3,1).contiguous().view(N*H*W, -1)
+            refLoss_idx = self.entropy_loss(ref_probs, gt_labels.squeeze(dim=1))
+        reconLoss_idx = torch.zeros_like(palLoss_idx)
+        if self.enhanced:
+            scalar = 1.0 if self.hint2regress else 5.0
+            reconLoss_idx = scalar * self._perceptual_loss(data['input_gray'], data['pred_color'], data['input_color'])
+            if self.with_grad:
+                gradient_loss = self._laplace_gradient(data['pred_color'], data['input_color'])
+                reconLoss_idx += gradient_loss
+        totalLoss_idx = palLoss_idx + refLoss_idx + reconLoss_idx
+        #print("loss terms:", palLoss_idx.item(), refLoss_idx.item(), reconLoss_idx.item())
+        return {'totalLoss':totalLoss_idx, 'palLoss':palLoss_idx, 'refLoss':refLoss_idx, 'recLoss':reconLoss_idx}
+
+
+def compute_affinity_pos_loss(prob_in, labxy_feat, pos_weight=0.003, kernel_size=16):
+    S = kernel_size
+    m = pos_weight
+    prob = prob_in.clone()
+    N,C,H,W = labxy_feat.shape
+    pooled_labxy = basic.poolfeat(labxy_feat, prob, kernel_size, kernel_size)
+    reconstr_feat = basic.upfeat(pooled_labxy, prob, kernel_size, kernel_size)
+    loss_map = reconstr_feat[:,:,:,:] - labxy_feat[:,:,:,:]
+    loss_feat = torch.norm(loss_map[:,:-2,:,:], p=2, dim=1).mean()
+    loss_pos = torch.norm(loss_map[:,-2:,:,:], p=2, dim=1).mean() * m / S
+    loss_affinity = loss_feat + loss_pos
+    return loss_affinity
+
+
+def l2_loss(y_input, y_target, weight_map=None):
+    if weight_map is None:
+        return F.mse_loss(y_input, y_target)
+    else:
+        diff_map = torch.mean(torch.abs(y_input-y_target), dim=1, keepdim=True)
+        batch_dev = torch.sum(diff_map*diff_map*weight_map, dim=(1,2,3)) / (eps+torch.sum(weight_map, dim=(1,2,3)))
+        return batch_dev.mean()
+
+
+def l1_loss(y_input, y_target, weight_map=None):
+    if weight_map is None:
+        return F.l1_loss(y_input, y_target)
+    else:
+        diff_map = torch.mean(torch.abs(y_input-y_target), dim=1, keepdim=True)
+        batch_dev = torch.sum(diff_map*weight_map, dim=(1,2,3)) / (eps+torch.sum(weight_map, dim=(1,2,3)))
+        return batch_dev.mean()
+
+
+def masked_l1_loss(y_input, y_target, outlier_mask):
+    one = torch.tensor([1.0]).cuda(y_input.get_device())
+    weight_map = torch.where(outlier_mask, one * 0.0, one * 1.0)
+    return l1_loss(y_input, y_target, weight_map)
+
+
+def huber_loss(y_input, y_target, delta=0.01):
+    mask = torch.zeros_like(y_input)
+    mann = torch.abs(y_input - y_target)
+    eucl = 0.5 * (mann**2)
+    mask[...] = mann < delta
+    loss = eucl * mask / delta + (mann - 0.5 * delta) * (1 - mask)
+    return torch.mean(loss)
+
+
+## Perceptual loss that uses a pretrained VGG network
+class VGG19Loss(nn.Module):
+    def __init__(self, feat_type='liu', gpu_no=0, is_ddp=False, requires_grad=False):
+        super(VGG19Loss, self).__init__()
+        os.environ['TORCH_HOME'] = '/apdcephfs/share_1290939/richardxia/Saved/Checkpoints/VGG19'
+        ## data requirement: (N,C,H,W) in RGB format, [0,1] range, and resolution >= 224x224
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]
+        self.feat_type = feat_type
+
+        vgg_model = torchvision.models.vgg19(pretrained=True)
+        ## AssertionError: DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient
+        '''
+        if is_ddp:
+            vgg_model = vgg_model.cuda(gpu_no)
+            vgg_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(vgg_model)
+            vgg_model = torch.nn.parallel.DistributedDataParallel(vgg_model, device_ids=[gpu_no], find_unused_parameters=True)
+        else:
+            vgg_model = vgg_model.cuda(gpu_no)
+        '''
+        vgg_model = vgg_model.cuda(gpu_no)
+        if self.feat_type == 'liu':
+            ## conv1_1, conv2_1, conv3_1, conv4_1, conv5_1
+            self.slice1 = nn.Sequential(*list(vgg_model.features)[:2]).eval()
+            self.slice2 = nn.Sequential(*list(vgg_model.features)[2:7]).eval()
+            self.slice3 = nn.Sequential(*list(vgg_model.features)[7:12]).eval()
+            self.slice4 = nn.Sequential(*list(vgg_model.features)[12:21]).eval()
+            self.slice5 = nn.Sequential(*list(vgg_model.features)[21:30]).eval()
+            self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
+        elif self.feat_type == 'lei':
+            ## conv1_2, conv2_2, conv3_2, conv4_2, conv5_2
+            self.slice1 = nn.Sequential(*list(vgg_model.features)[:4]).eval()
+            self.slice2 = nn.Sequential(*list(vgg_model.features)[4:9]).eval()
+            self.slice3 = nn.Sequential(*list(vgg_model.features)[9:14]).eval()
+            self.slice4 = nn.Sequential(*list(vgg_model.features)[14:23]).eval()
+            self.slice5 = nn.Sequential(*list(vgg_model.features)[23:32]).eval()
+            self.weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10.0/1.5]
+        else:
+            ## maxpool after conv4_4
+            self.featureExactor = nn.Sequential(*list(vgg_model.features)[:28]).eval()
+        '''
+        for x in range(2):
+            self.slice1.add_module(str(x), pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), pretrained_features[x])
+        '''
+        self.criterion = nn.L1Loss()
+
+        ## fixed parameters
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+        self.eval()
+        print('[*] VGG19Loss init!')
+
+    def normalize(self, tensor):
+        tensor = tensor.clone()
+        mean = torch.as_tensor(self.mean, dtype=torch.float32, device=tensor.device)
+        std = torch.as_tensor(self.std, dtype=torch.float32, device=tensor.device)
+        tensor.sub_(mean[None, :, None, None]).div_(std[None, :, None, None])
+        return tensor
+
+    def forward(self, x, y):
+        norm_x, norm_y = self.normalize(x), self.normalize(y)
+        ## feature extract
+        if self.feat_type == 'liu' or self.feat_type == 'lei':
+            x_relu1, y_relu1 = self.slice1(norm_x), self.slice1(norm_y)
+            x_relu2, y_relu2 = self.slice2(x_relu1), self.slice2(y_relu1)
+            x_relu3, y_relu3 = self.slice3(x_relu2), self.slice3(y_relu2)
+            x_relu4, y_relu4 = self.slice4(x_relu3), self.slice4(y_relu3)
+            x_relu5, y_relu5 = self.slice5(x_relu4), self.slice5(y_relu4)
+            x_vgg = [x_relu1, x_relu2, x_relu3, x_relu4, x_relu5]
+            y_vgg = [y_relu1, y_relu2, y_relu3, y_relu4, y_relu5]
+            loss = 0    
+            for i in range(len(x_vgg)):
+                loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
+        else:
+            x_vgg, y_vgg = self.featureExactor(norm_x), self.featureExactor(norm_y)
+            loss = self.criterion(x_vgg, y_vgg.detach())
+        return loss