Update networks.py

networks.py

@@ -1,13 +1,11 @@
-# coding=utf-8
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.nn import init
 from torchvision import models
 import os
-
 import numpy as np
 
-
 def weights_init_normal(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
@@ -18,7 +16,6 @@ def weights_init_normal(m):
         init.normal_(m.weight.data, 1.0, 0.02)
         init.constant_(m.bias.data, 0.0)
 
-
 def weights_init_xavier(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
@@ -29,7 +26,6 @@ def weights_init_xavier(m):
         init.normal_(m.weight.data, 1.0, 0.02)
         init.constant_(m.bias.data, 0.0)
 
-
 def weights_init_kaiming(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
@@ -40,7 +36,6 @@ def weights_init_kaiming(m):
         init.normal_(m.weight.data, 1.0, 0.02)
         init.constant_(m.bias.data, 0.0)
 
-
 def init_weights(net, init_type='normal'):
     print('initialization method [%s]' % init_type)
     if init_type == 'normal':
@@ -53,7 +48,6 @@ def init_weights(net, init_type='normal'):
         raise NotImplementedError(
             'initialization method [%s] is not implemented' % init_type)
 
-
 class FeatureExtraction(nn.Module):
     def __init__(self, input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_dropout=False):
         super(FeatureExtraction, self).__init__()
@@ -78,7 +72,6 @@ class FeatureExtraction(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-
 class FeatureL2Norm(torch.nn.Module):
     def __init__(self):
         super(FeatureL2Norm, self).__init__()
@@ -89,7 +82,6 @@ class FeatureL2Norm(torch.nn.Module):
                          epsilon, 0.5).unsqueeze(1).expand_as(feature)
         return torch.div(feature, norm)
 
-
 class FeatureCorrelation(nn.Module):
     def __init__(self):
         super(FeatureCorrelation, self).__init__()
@@ -105,9 +97,8 @@ class FeatureCorrelation(nn.Module):
             b, h, w, h*w).transpose(2, 3).transpose(1, 2)
         return correlation_tensor
 
-
 class FeatureRegression(nn.Module):
-    def __init__(self, input_nc=512, output_dim=6, use_cuda=True):
+    def __init__(self, input_nc=512, output_dim=6):
         super(FeatureRegression, self).__init__()
         self.conv = nn.Sequential(
             nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1),
@@ -125,10 +116,6 @@ class FeatureRegression(nn.Module):
         )
         self.linear = nn.Linear(64 * 4 * 3, output_dim)
         self.tanh = nn.Tanh()
-        if use_cuda:
-            self.conv.cuda()
-            self.linear.cuda()
-            self.tanh.cuda()
 
     def forward(self, x):
         x = self.conv(x)
@@ -137,7 +124,6 @@ class FeatureRegression(nn.Module):
         x = self.tanh(x)
         return x
 
-
 class AffineGridGen(nn.Module):
     def __init__(self, out_h=256, out_w=192, out_ch=3):
         super(AffineGridGen, self).__init__()
@@ -152,13 +138,11 @@ class AffineGridGen(nn.Module):
             (batch_size, self.out_ch, self.out_h, self.out_w))
         return F.affine_grid(theta, out_size)
 
-
 class TpsGridGen(nn.Module):
-    def __init__(self, out_h=256, out_w=192, use_regular_grid=True, grid_size=3, reg_factor=0, use_cuda=True):
+    def __init__(self, out_h=256, out_w=192, use_regular_grid=True, grid_size=3, reg_factor=0):
         super(TpsGridGen, self).__init__()
         self.out_h, self.out_w = out_h, out_w
         self.reg_factor = reg_factor
-        self.use_cuda = use_cuda
 
         # create grid in numpy
         self.grid = np.zeros([self.out_h, self.out_w, 3], dtype=np.float32)
@@ -168,9 +152,6 @@ class TpsGridGen(nn.Module):
         # grid_X,grid_Y: size [1,H,W,1,1]
         self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3)
         self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3)
-        if use_cuda:
-            self.grid_X = self.grid_X.cuda()
-            self.grid_Y = self.grid_Y.cuda()
 
         # initialize regular grid for control points P_i
         if use_regular_grid:
@@ -188,11 +169,6 @@ class TpsGridGen(nn.Module):
                 3).unsqueeze(4).transpose(0, 4)
             self.P_Y = P_Y.unsqueeze(2).unsqueeze(
                 3).unsqueeze(4).transpose(0, 4)
-            if use_cuda:
-                self.P_X = self.P_X.cuda()
-                self.P_Y = self.P_Y.cuda()
-                self.P_X_base = self.P_X_base.cuda()
-                self.P_Y_base = self.P_Y_base.cuda()
 
     def forward(self, theta):
         warped_grid = self.apply_transformation(
@@ -217,8 +193,6 @@ class TpsGridGen(nn.Module):
         L = torch.cat((torch.cat((K, P), 1), torch.cat(
             (P.transpose(0, 1), Z), 1)), 0)
         Li = torch.inverse(L)
-        if self.use_cuda:
-            Li = Li.cuda()
         return Li
 
     def apply_transformation(self, theta, points):
@@ -315,8 +289,6 @@ class TpsGridGen(nn.Module):
 # |num_downs|: number of downsamplings in UNet. For example,
 # if |num_downs| == 7, image of size 128x128 will become of size 1x1
 # at the bottleneck
-
-
 class UnetGenerator(nn.Module):
     def __init__(self, input_nc, output_nc, num_downs, ngf=64,
                  norm_layer=nn.BatchNorm2d, use_dropout=False):
@@ -341,7 +313,6 @@ class UnetGenerator(nn.Module):
     def forward(self, input):
         return self.model(input)
 
-
 # Defines the submodule with skip connection.
 # X -------------------identity---------------------- X
 #   |-- downsampling -- |submodule| -- upsampling --|
@@ -395,7 +366,6 @@ class UnetSkipConnectionBlock(nn.Module):
         else:
             return torch.cat([x, self.model(x)], 1)
 
-
 class Vgg19(nn.Module):
     def __init__(self, requires_grad=False):
         super(Vgg19, self).__init__()
@@ -428,12 +398,10 @@ class Vgg19(nn.Module):
         out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
         return out
 
-
 class VGGLoss(nn.Module):
     def __init__(self, layids=None):
         super(VGGLoss, self).__init__()
         self.vgg = Vgg19()
-        self.vgg.cuda()
         self.criterion = nn.L1Loss()
         self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
         self.layids = layids
@@ -448,7 +416,6 @@ class VGGLoss(nn.Module):
                 self.criterion(x_vgg[i], y_vgg[i].detach())
         return loss
 
-
 class DT(nn.Module):
     def __init__(self):
         super(DT, self).__init__()
@@ -457,17 +424,15 @@ class DT(nn.Module):
         dt = torch.abs(x1 - x2)
         return dt
 
-
 class DT2(nn.Module):
     def __init__(self):
-        super(DT, self).__init__()
+        super(DT2, self).__init__()
 
     def forward(self, x1, y1, x2, y2):
         dt = torch.sqrt(torch.mul(x1 - x2, x1 - x2) +
                         torch.mul(y1 - y2, y1 - y2))
         return dt
 
-
 class GicLoss(nn.Module):
     def __init__(self, opt):
         super(GicLoss, self).__init__()
@@ -496,7 +461,6 @@ class GicLoss(nn.Module):
 
         return torch.sum(torch.abs(dtleft - dtright) + torch.abs(dtup - dtdown))
 
-
 class GMM(nn.Module):
     """ Geometric Matching Module
     """
@@ -510,9 +474,9 @@ class GMM(nn.Module):
         self.l2norm = FeatureL2Norm()
         self.correlation = FeatureCorrelation()
         self.regression = FeatureRegression(
-            input_nc=192, output_dim=2*opt.grid_size**2, use_cuda=True)
+            input_nc=192, output_dim=2*opt.grid_size**2)
         self.gridGen = TpsGridGen(
-            opt.fine_height, opt.fine_width, use_cuda=True, grid_size=opt.grid_size)
+            opt.fine_height, opt.fine_width, grid_size=opt.grid_size)
 
     def forward(self, inputA, inputB):
         featureA = self.extractionA(inputA)
@@ -525,17 +489,12 @@ class GMM(nn.Module):
         grid = self.gridGen(theta)
         return grid, theta
 
-
 def save_checkpoint(model, save_path):
     if not os.path.exists(os.path.dirname(save_path)):
         os.makedirs(os.path.dirname(save_path))
-
-    torch.save(model.cpu().state_dict(), save_path)
-    model.cuda()
-
+    torch.save(model.state_dict(), save_path)
 
 def load_checkpoint(model, checkpoint_path):
     if not os.path.exists(checkpoint_path):
         return
-    model.load_state_dict(torch.load(checkpoint_path))
-    model.cuda()
+    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))