def create_conv_layers(self):
dec_convs = []
nf = self.nf_dcgan
dec_convs.append(
dcgan_upconv(self.in_channels[0] + self.skip_n_channels[-1], nf * 4))
dec_convs.append(
dcgan_upconv(nf * 4 + self.skip_n_channels[-2], nf * 2))
dec_convs.append(
dcgan_upconv(nf * 2 + self.skip_n_channels[-3], nf))
if self.img_size[0] == 64:
dec_convs.append(
dcgan_upconv(nf + self.skip_n_channels[-4], nf))
self.dec_convs = nn.ModuleList(dec_convs)
elif self.img_size[0] == 32:
self.dec_convs = nn.ModuleList(dec_convs)
else:
raise NotImplementedError
nc = self.img_size[-1]
self.last_upsample = None
self.last_conv = nn.Sequential(
nn.ConvTranspose2d(nf + self.skip_n_channels[0], nc, 4, 2, 1))
for module in self.last_conv:
_init_layer(module)
def create_conv_layers(self, ic, ks, nm, nl, st):
# 32 x 32 / 64 x 64
self.c64 = None
nf = self.nf_vgg
if self.img_size[0] == 64:
self.c64 = nn.Sequential(
vgg_layer(self.ic[0], self.ic[0]),
vgg_layer(self.ic[0], self.ic[0]),
)
self.c0 = nn.Sequential(
vgg_layer(self.ic[0], nf),
vgg_layer(nf, nf),
)
self.c1 = nn.Sequential(
vgg_layer(nf, nf * 2),
vgg_layer(nf * 2, nf * 2),
vgg_layer(nf * 2, nf * 2),
)
# / # 4 x 4 - depending on the input image resolution
self.c2 = nn.Sequential(
vgg_layer(nf * 2, nf * 4),
vgg_layer(nf * 4, nf * 4),
vgg_layer(nf * 4, nf * 4),
)
self.c3 = nn.Sequential(
nn.Conv2d(nf * 4, self.oc[2], 3, 1, 1),
nn.BatchNorm2d(self.oc[2]),
)
for module in self.c3:
_init_layer(module)
self.mp = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
def create_conv_layers(self):
dec_convs = []
nf = self.nf_vgg
self.upc1 = nn.Sequential(
nn.ConvTranspose2d(self.in_channels[0] + self.skip_n_channels[-1], nf*4, 3, 1, 1),
nn.BatchNorm2d(nf * 4),
nn.LeakyReLU(0.2, inplace=True)
)
dec_convs.append(
nn.Sequential(
vgg_layer(nf*4 + self.skip_n_channels[-2], nf*4),
vgg_layer(nf*4, nf*4),
vgg_layer(nf*4, nf*2),
)
) # 8 x 8
dec_convs.append(
nn.Sequential(
vgg_layer(nf*2 + self.skip_n_channels[-3], nf*2),
vgg_layer(nf*2, nf*2),
vgg_layer(nf*2, nf*2),
)
) # 16 x 16
if self.img_size[0] == 64:
dec_convs.append(
nn.Sequential(
vgg_layer(nf*2 + self.skip_n_channels[-4], nf*2),
vgg_layer(nf*2, nf),
)
) # 32 x 32
self.dec_convs = nn.ModuleList(dec_convs)
out_channels = self.img_size[-1]
self.last_conv = nn.Sequential(
vgg_layer(nf + self.skip_n_channels[0], nf),
nn.ConvTranspose2d(nf, out_channels, 3, 1, 1),
) # 64 x 64
elif self.img_size[0] == 32:
self.dec_convs = nn.ModuleList(dec_convs)
out_channels = self.img_size[-1]
self.last_conv = nn.Sequential(
vgg_layer(nf*2 + self.skip_n_channels[0], nf),
nn.ConvTranspose2d(nf, out_channels, 3, 1, 1),
) # 32 x 32
else:
raise NotImplementedError
self.up = nn.UpsamplingNearest2d(scale_factor=2)
for module in self.upc1:
_init_layer(module)
for module in self.last_conv:
_init_layer(module)