def __init__(self, pose_size, start_channel_dim, image_channels):
super().__init__(pose_size, start_channel_dim, image_channels)
# Transition blockss
self.to_rgb_new = WSConv2d(start_channel_dim, self.image_channels, 1,
0)
self.to_rgb_old = WSConv2d(start_channel_dim, self.image_channels, 1,
0)
self.core_blocks_down = nn.ModuleList(
[UnetDownSamplingBlock(start_channel_dim, start_channel_dim)])
self.core_blocks_up = nn.ModuleList([
nn.Sequential(
conv_bn_relu(start_channel_dim + self.num_poses + 32,
start_channel_dim, 1, 0),
UnetUpsamplingBlock(start_channel_dim, start_channel_dim))
])
self.new_up = nn.Sequential()
self.old_up = nn.Sequential()
self.new_down = nn.Sequential()
self.from_rgb_new = conv_bn_relu(self.image_channels,
start_channel_dim, 1, 0)
self.from_rgb_old = conv_bn_relu(self.image_channels,
start_channel_dim, 1, 0)
self.upsampling = UpSamplingBlock()
self.downsampling = nn.AvgPool2d(2)
def extend(self):
output_dim = self.transition_channels[self.transition_step]
print("extending G", output_dim)
# Downsampling module
self.from_rgb_old = nn.Sequential(
nn.AvgPool2d([2, 2]),
self.from_rgb_new
)
if self.transition_step == 0:
core_block_up = nn.Sequential(
self.core_blocks_up[0],
UpSamplingBlock()
)
self.core_blocks_up = nn.ModuleList([core_block_up])
else:
core_blocks_down = nn.ModuleList()
core_blocks_down.append(self.new_down)
first = [nn.AvgPool2d(2)] + \
list(self.core_blocks_down[0].children())
core_blocks_down.append(nn.Sequential(*first))
core_blocks_down.extend(self.core_blocks_down[1:])
self.core_blocks_down = core_blocks_down
new_up_blocks = list(self.new_up.children()) + [UpSamplingBlock()]
self.new_up = nn.Sequential(*new_up_blocks)
self.core_blocks_up.append(self.new_up)
self.from_rgb_new = conv_bn_relu(self.image_channels, output_dim, 1, 0)
self.new_down = nn.Sequential(
UnetDownSamplingBlock(output_dim, self.prev_channel_extension)
)
self.new_down = self.new_down
# Upsampling modules
self.to_rgb_old = self.to_rgb_new
self.to_rgb_new = WSConv2d(output_dim, self.image_channels, 1, 0)
self.new_up = nn.Sequential(
conv_bn_relu(self.prev_channel_extension*2 +
self.num_poses, self.prev_channel_extension, 1, 0),
UnetUpsamplingBlock(self.prev_channel_extension, output_dim)
)
super().extend()
def conv_module_bn(dim_in, dim_out, kernel_size, padding):
return nn.Sequential(WSConv2d(dim_in, dim_out, kernel_size, padding),
nn.LeakyReLU(negative_slope=.2))
def conv_bn_relu(in_dim, out_dim, kernel_size, padding=0):
return nn.Sequential(WSConv2d(in_dim, out_dim, kernel_size, padding),
nn.LeakyReLU(negative_slope=.2),
PixelwiseNormalization())