-
Notifications
You must be signed in to change notification settings - Fork 0
/
Generators.py
executable file
·155 lines (134 loc) · 6.01 KB
/
Generators.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
#!/usr/bin/python
'''
Defines the Encoders and Decoder that make up the
Enc+Dec/Enc+2Dec model defined in DLoc
Idea and code originally from fstin Johnson's architecture.
https://github.com/jcjohnson/fast-neural-style/
Code base expanded from pix2pix Phillip Isola's implementation
https://github.com/phillipi/pix2pix
You can add your costum network building blocks here to test various other architectures.
'''
import torch
import torch.nn as nn
import numpy as np
import functools
from params import opt_exp
# Tha base Encoder function defined for the DLoc's Encoder
class ResnetEncoder(nn.Module):
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='zero'):
assert(n_blocks >= 0)
super(ResnetEncoder, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = [nn.Conv2d(input_nc, input_nc, kernel_size=7, padding=3,
bias=use_bias),
norm_layer(input_nc),
nn.Tanh()]
model += [nn.Conv2d(input_nc, ngf, kernel_size=7, padding=3,
bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)]
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
stride=2, padding=1, bias=use_bias),
norm_layer(ngf * mult * 2),
nn.ReLU(True)]
mult = 2**n_downsampling
for i in range(n_blocks):
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
self.model = nn.Sequential(*model)
def forward(self, input):
return self.model(input)
# Tha base Decoder function defined for the DLoc's Decoders.
# Depending upon the ModelADT wraper around the decoder,
# the decoder would either be a Location Decoder or a Consistency decoder.
# For more details refer to ModelADT.py wrapper implementation and params.py
class ResnetDecoder(nn.Module):
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=9, padding_type='zero', encoder_blocks=6):
assert(n_blocks >= 0)
super(ResnetDecoder, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = []
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
mult = 2**n_downsampling
for i in range(n_blocks):
if i <= encoder_blocks:
continue
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
if n_blocks == 9:
model += [nn.Conv2d(ngf * mult, 32, kernel_size=3,
stride=1, padding=1, bias=use_bias),
nn.ReLU()]
model += [nn.Flatten()]
linear_shape = int(np.prod((opt_exp.batch_size, 32, 41, 91))/opt_exp.batch_size)
model += [nn.Linear(linear_shape, 256), nn.ReLU()]
model += [nn.Linear(256, 256), nn.ReLU()]
model += [nn.Linear(256, 2)]
model += [nn.Tanh()]
else:
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
kernel_size=3, stride=2,
padding=1, output_padding=0,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)]
# model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=[3,3])]
model += [nn.Sigmoid()]
self.model = nn.Sequential(*model)
def forward(self, input):
#print("decoder.input = ", input.shape)
return self.model(input)
# Define a resnet block
class ResnetBlock(nn.Module):
def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
super(ResnetBlock, self).__init__()
self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim),
nn.ReLU(True)]
if use_dropout:
conv_block += [nn.Dropout(0.25)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim)]
return nn.Sequential(*conv_block)
def forward(self, x):
out = x + self.conv_block(x)
return out