FusionGAN.zip

import torch import torch.nn as nn import torch.utils.model_zoo as model_zoo import math from torch.autograd import Variable def conv2d(in_channels, out_channels, kernel_size = 3, padding = 1): return nn.Conv2d(in_channels, out_channels, kernel_size = kernel_size, padding = padding) def deconv2d(in_channels, out_channels, kernel_size = 3, padding = 1): return nn.ConvTranspose2d(in_channels, out_channels, kernel_size = kernel_size, padding = padding) def relu(inplace = True): # Change to True? return nn.ReLU(inplace) def maxpool2d(): return nn.MaxPool2d(2) def make_conv_layers(cfg): layers = [] in_channels = 3 for v in cfg: if v == 'M': layers += [maxpool2d()] else: conv = conv2d(in_channels, v) layers += [conv, relu(inplace=True)] in_channels = v return nn.Sequential(*layers) def make_deconv_layers(cfg): layers = [] in_channels = 512 for v in cfg: if v == 'U': layers += [nn.Upsample(scale_factor=2)] else: deconv = deconv2d(in_channels, v) layers += [deconv] in_channels = v return nn.Sequential(*layers) cfg = { 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512], 'D': [512, 512, 512, 'U', 512, 512, 512, 'U', 256, 256, 256, 'U', 128, 128, 'U', 64, 64] } def encoder(): return make_conv_layers(cfg['E']) def decoder(): return make_deconv_layers(cfg['D']) class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.encoder = encoder() self.decoder = decoder() self.mymodules = nn.ModuleList([ deconv2d(64,1,kernel_size=1, padding = 0), nn.Sigmoid() ]) def forward(self,x): # #print('Input x', x.size()) x = self.encoder(x) #print('After encoder = ', x.size()) x = self.decoder(x) #print('After decoder = ', x.size()) x = self.mymodules[0](x) x = self.mymodules[1](x) #print('Final size = ', x.size()) return x class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() self.convs = nn.Sequential( # [-1, 4, 256,192] conv2d(4, 3, kernel_size=1), relu(), conv2d(3, 32, kernel_size=3), # [-1, 32, 256, 192] relu(), maxpool2d(), conv2d(32, 64, kernel_size=3), # [-1, 64, 128, 96] relu(), conv2d(64, 64, kernel_size=3), # [-1, 64, 128, 96] relu(), maxpool2d(), # [-1,64,64,48] conv2d(64, 64, kernel_size=3), # [-1,64,64,48] relu(), conv2d(64, 64, kernel_size=3), relu(), maxpool2d(), # [-1,64,32,24] ) self.mymodules = nn.ModuleList([ nn.Sequential(nn.Linear(64 * 32 * 24, 100), nn.Tanh()), nn.Sequential(nn.Linear(100, 2), nn.Tanh()), nn.Sequential(nn.Linear(2, 1), nn.Sigmoid()) ]) # self._initialize_weights() def forward(self, x): x = self.convs(x) x = x.view(-1, self.num_flat_features(x)) x = self.mymodules[0](x) x = self.mymodules[1](x) x = self.mymodules[2](x) return x def num_flat_features(self, x): size = x.size()[1:] # all dimensions except the batch dimension num_features = 1 for s in size: num_features *= s return num_features def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels # print(m.weight.data.shape) # print('old conv2d layer!') # print(m.weight.data.min()) # print(m.weight.data.max()) m.weight.data.normal_(0, math.sqrt(2. / n)) # print('new conv2d layer!') # print(m.weight.data.min()) # print(m.weight.data.max()) if m.bias is not None: m.bias.data.zero_() elif isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01) m.bias.data.zero_()