95 lines
4.1 KiB
Python
95 lines
4.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from model.base.module import ResidualBlock, ReverseConv2dBlock, Conv2dBlock
|
|
|
|
|
|
class StyleEncoder(nn.Module):
|
|
def __init__(self, in_channels, style_dim, num_conv, end_size=(4, 4), base_channels=64,
|
|
norm_type="IN", padding_mode='reflect', activation_type="LeakyReLU"):
|
|
super().__init__()
|
|
sequence = [Conv2dBlock(
|
|
in_channels, base_channels, kernel_size=3, stride=1, padding=1, padding_mode=padding_mode,
|
|
activation_type=activation_type, norm_type=norm_type
|
|
)]
|
|
multiple_now = 0
|
|
max_multiple = 3
|
|
for i in range(1, num_conv + 1):
|
|
multiple_prev = multiple_now
|
|
multiple_now = min(2 ** i, 2 ** max_multiple)
|
|
sequence.append(Conv2dBlock(
|
|
multiple_prev * base_channels, multiple_now * base_channels,
|
|
kernel_size=3, stride=2, padding=1, padding_mode=padding_mode,
|
|
activation_type=activation_type, norm_type=norm_type
|
|
))
|
|
self.sequence = nn.Sequential(*sequence)
|
|
self.fc_avg = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)
|
|
self.fc_var = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)
|
|
|
|
def forward(self, x):
|
|
x = self.sequence(x)
|
|
x = x.view(x.size(0), -1)
|
|
return self.fc_avg(x), self.fc_var(x)
|
|
|
|
|
|
class SPADEGenerator(nn.Module):
|
|
def __init__(self, in_channels, out_channels, num_blocks, use_vae, num_z_dim, start_size=(4, 4), base_channels=64,
|
|
padding_mode='reflect', activation_type="LeakyReLU"):
|
|
super().__init__()
|
|
self.sx, self.sy = start_size
|
|
self.use_vae = use_vae
|
|
self.num_z_dim = num_z_dim
|
|
if use_vae:
|
|
self.input_converter = nn.Linear(num_z_dim, 16 * base_channels * self.sx * self.sy)
|
|
else:
|
|
self.input_converter = nn.Conv2d(in_channels, 16 * base_channels, kernel_size=3, padding=1)
|
|
|
|
sequence = []
|
|
|
|
multiple_now = 16
|
|
for i in range(num_blocks - 1, -1, -1):
|
|
multiple_prev = multiple_now
|
|
multiple_now = min(2 ** i, 2 ** 4)
|
|
if i != num_blocks - 1:
|
|
sequence.append(nn.Upsample(scale_factor=2))
|
|
sequence.append(ResidualBlock(
|
|
base_channels * multiple_prev,
|
|
out_channels=base_channels * multiple_now,
|
|
padding_mode=padding_mode,
|
|
activation_type=activation_type,
|
|
norm_type="SPADE",
|
|
pre_activation=True,
|
|
additional_norm_kwargs=dict(
|
|
condition_in_channels=in_channels, base_channels=128, base_norm_type="BN",
|
|
activation_type="ReLU", padding_mode="zeros", gamma_bias=1.0
|
|
)
|
|
))
|
|
self.sequence = nn.Sequential(*sequence)
|
|
self.output_converter = nn.Sequential(
|
|
ReverseConv2dBlock(base_channels, out_channels, kernel_size=3, stride=1, padding=1,
|
|
padding_mode=padding_mode, activation_type=activation_type, norm_type="NONE"),
|
|
nn.Tanh()
|
|
)
|
|
|
|
def forward(self, seg, z=None):
|
|
if self.use_vae:
|
|
if z is None:
|
|
z = torch.randn(seg.size(0), self.num_z_dim, device=seg.device)
|
|
x = self.input_converter(z).view(seg.size(0), -1, self.sx, self.sy)
|
|
else:
|
|
x = self.input_converter(F.interpolate(seg, size=(self.sx, self.sy)))
|
|
for blk in self.sequence:
|
|
if isinstance(blk, ResidualBlock):
|
|
downsampling_seg = F.interpolate(seg, size=x.size()[2:], mode='nearest')
|
|
blk.conv1.normalization.set_condition_image(downsampling_seg)
|
|
blk.conv2.normalization.set_condition_image(downsampling_seg)
|
|
if blk.learn_skip_connection:
|
|
blk.res_conv.normalization.set_condition_image(downsampling_seg)
|
|
x = blk(x)
|
|
return self.output_converter(x)
|
|
|
|
if __name__ == '__main__':
|
|
g = SPADEGenerator(3, 3, 7, False, 256)
|
|
print(g)
|
|
print(g(torch.randn(2, 3, 256, 256)).size()) |