Delete train_log/loss.py

train_log/loss.py

@@ -1,128 +0,0 @@
-import torch
-import numpy as np
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.models as models
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-class EPE(nn.Module):
-    def __init__(self):
-        super(EPE, self).__init__()
-
-    def forward(self, flow, gt, loss_mask):
-        loss_map = (flow - gt.detach()) ** 2
-        loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
-        return (loss_map * loss_mask)
-
-
-class Ternary(nn.Module):
-    def __init__(self):
-        super(Ternary, self).__init__()
-        patch_size = 7
-        out_channels = patch_size * patch_size
-        self.w = np.eye(out_channels).reshape(
-            (patch_size, patch_size, 1, out_channels))
-        self.w = np.transpose(self.w, (3, 2, 0, 1))
-        self.w = torch.tensor(self.w).float().to(device)
-
-    def transform(self, img):
-        patches = F.conv2d(img, self.w, padding=3, bias=None)
-        transf = patches - img
-        transf_norm = transf / torch.sqrt(0.81 + transf**2)
-        return transf_norm
-
-    def rgb2gray(self, rgb):
-        r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
-        gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
-        return gray
-
-    def hamming(self, t1, t2):
-        dist = (t1 - t2) ** 2
-        dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
-        return dist_norm
-
-    def valid_mask(self, t, padding):
-        n, _, h, w = t.size()
-        inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
-        mask = F.pad(inner, [padding] * 4)
-        return mask
-
-    def forward(self, img0, img1):
-        img0 = self.transform(self.rgb2gray(img0))
-        img1 = self.transform(self.rgb2gray(img1))
-        return self.hamming(img0, img1) * self.valid_mask(img0, 1)
-
-
-class SOBEL(nn.Module):
-    def __init__(self):
-        super(SOBEL, self).__init__()
-        self.kernelX = torch.tensor([
-            [1, 0, -1],
-            [2, 0, -2],
-            [1, 0, -1],
-        ]).float()
-        self.kernelY = self.kernelX.clone().T
-        self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
-        self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
-
-    def forward(self, pred, gt):
-        N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
-        img_stack = torch.cat(
-            [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0)
-        sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
-        sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
-        pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:]
-        pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:]
-
-        L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y)
-        loss = (L1X+L1Y)
-        return loss
-
-class MeanShift(nn.Conv2d):
-    def __init__(self, data_mean, data_std, data_range=1, norm=True):
-        c = len(data_mean)
-        super(MeanShift, self).__init__(c, c, kernel_size=1)
-        std = torch.Tensor(data_std)
-        self.weight.data = torch.eye(c).view(c, c, 1, 1)
-        if norm:
-            self.weight.data.div_(std.view(c, 1, 1, 1))
-            self.bias.data = -1 * data_range * torch.Tensor(data_mean)
-            self.bias.data.div_(std)
-        else:
-            self.weight.data.mul_(std.view(c, 1, 1, 1))
-            self.bias.data = data_range * torch.Tensor(data_mean)
-        self.requires_grad = False
-            
-class VGGPerceptualLoss(torch.nn.Module):
-    def __init__(self, rank=0):
-        super(VGGPerceptualLoss, self).__init__()
-        blocks = []
-        pretrained = True
-        self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
-        self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def forward(self, X, Y, indices=None):
-        X = self.normalize(X)
-        Y = self.normalize(Y)
-        indices = [2, 7, 12, 21, 30]
-        weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5]
-        k = 0
-        loss = 0
-        for i in range(indices[-1]):
-            X = self.vgg_pretrained_features[i](X)
-            Y = self.vgg_pretrained_features[i](Y)
-            if (i+1) in indices:
-                loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
-                k += 1
-        return loss
-
-if __name__ == '__main__':
-    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
-    img1 = torch.tensor(np.random.normal(
-        0, 1, (3, 3, 256, 256))).float().to(device)
-    ternary_loss = Ternary()
-    print(ternary_loss(img0, img1).shape)