Update external_function.py

Author: PangzeCheung

models/external_function.py

@@ -149,127 +149,6 @@ def reduce_sum(x, axis=None, keepdim=False):
     return x
 
 
-class CorrelationLoss(nn.Module):
-    r"""
-    Perceptual loss, VGG-based
-    https://arxiv.org/abs/1603.08155
-    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
-    """
-
-    def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]):
-        super(CorrelationLoss, self).__init__()
-        self.add_module('vgg', VGG19())
-        self.criterion = torch.nn.L1Loss()
-        self.weights = weights
-        self.softmax = nn.Softmax(dim=-1)
-
-    def compute_gram(self, x):
-        b, ch, h, w = x.size()
-        f = x.view(b, ch, w * h)
-        f_T = f.transpose(1, 2)
-        G = f.bmm(f_T) / (h * w * ch)
-        return G
-
-    def comput_correlation(self, x_vgg, y_vgg, source_vgg):
-        input_shape = list(x_vgg.shape)
-
-        x_query = x_vgg.view(input_shape[0], -1, input_shape[2] * input_shape[3]).permute(0, 2,
-                                                                                                     1)  # B X N X C
-
-        y_query = y_vgg.view(input_shape[0], -1, input_shape[2] * input_shape[3]).permute(0, 2,
-                                                                                                     1)  # B X N X C
-
-        source_key = source_vgg.view(input_shape[0], -1, input_shape[2] * input_shape[3]).permute(0, 2,
-                                                                                                             1)  # B X N X C
-
-        escape_NaN = torch.FloatTensor([1e-4])
-        escape_NaN = escape_NaN.cuda()
-
-        max_x = torch.max(torch.sqrt(reduce_sum(torch.pow(x_query, 2),
-                                                axis=[2],
-                                                keepdim=True)), escape_NaN)
-        max_y = torch.max(torch.sqrt(reduce_sum(torch.pow(y_query, 2),
-                                                axis=[2],
-                                                keepdim=True)), escape_NaN)
-        max_source = torch.max(torch.sqrt(reduce_sum(torch.pow(source_key, 2),
-                                                     axis=[2],
-                                                     keepdim=True)), escape_NaN)
-
-        x_query = x_query / max_x
-        y_query = y_query / max_y
-        source_key = source_key / max_source
-
-        source_key = source_key.permute(0, 2, 1)
-
-        x_energy = torch.bmm(x_query, source_key)  # transpose check
-        x_attention = self.softmax(x_energy * 10)  # BX (N) X (N)
-
-        y_energy = torch.bmm(y_query, source_key)  # transpose check
-        y_attention = self.softmax(y_energy * 10)  # BX (N) X (N)
-        return x_attention, y_attention
-
-    def __call__(self, x, y, source):
-        # Compute features
-        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
-
-        content_loss = 0.0
-        content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1'])
-        content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1'])
-        content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1'])
-        content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1'])
-        content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])
-
-        # Compute loss
-        style_loss = 0.0
-        style_loss += self.criterion(self.compute_gram(x_vgg['relu2_2']), self.compute_gram(y_vgg['relu2_2']))
-        style_loss += self.criterion(self.compute_gram(x_vgg['relu3_4']), self.compute_gram(y_vgg['relu3_4']))
-        style_loss += self.criterion(self.compute_gram(x_vgg['relu4_4']), self.compute_gram(y_vgg['relu4_4']))
-        style_loss += self.criterion(self.compute_gram(x_vgg['relu5_2']), self.compute_gram(y_vgg['relu5_2']))
-
-        correlation_loss = 0
-
-        source_vgg = self.vgg(source)
-
-        x_attention_5, y_attention_5 = self.comput_correlation(x_vgg['relu5_1'], y_vgg['relu5_1'], source_vgg['relu5_1'])
-
-        correlation_loss += self.criterion(x_attention_5, y_attention_5)
-
-        x_attention_4, y_attention_4 = self.comput_correlation(x_vgg['relu4_1'], y_vgg['relu4_1'],
-                                                               source_vgg['relu4_1'])
-
-        correlation_loss += self.criterion(x_attention_4, y_attention_4)
-
-        return content_loss, style_loss, correlation_loss
-
-
-class DomainAdaptionLoss(nn.Module):
-    r"""
-    Perceptual loss, VGG-based
-    https://arxiv.org/abs/1603.08155
-    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
-    """
-
-    def __init__(self):
-        super(DomainAdaptionLoss, self).__init__()
-        self.criterion = torch.nn.L1Loss()
-
-    def compute_gram(self, x):
-        b, ch, h, w = x.size()
-        f = x.view(b, ch, w * h)
-        f_T = f.transpose(1, 2)
-        G = f.bmm(f_T) / (h * w * ch)
-        return G
-
-    def __call__(self, x, y):
-        # Compute features
-        content_loss = self.criterion(x, y)
-
-        # Compute loss
-        style_loss = self.criterion(self.compute_gram(x), self.compute_gram(y))
-
-        return content_loss, style_loss
-
-
 ####################################################################################################
 # neural style transform loss from neural_style_tutorial of pytorch
 ####################################################################################################
@@ -460,4 +339,4 @@ def forward(self, x):
             'relu5_3': relu5_3,
             'relu5_4': relu5_4,
         }
-        return out
+        return out