| | ---
|
| | license: mit
|
| | ---
|
| | Model weight for Fast Style Transfer
|
| |
|
| | ```
|
| |
|
| | class TransformerNetwork(nn.Module):
|
| | """Feedforward Transformation Network without Tanh
|
| | reference: https://arxiv.org/abs/1603.08155
|
| | exact architecture: https://cs.stanford.edu/people/jcjohns/papers/fast-style/fast-style-supp.pdf
|
| | """
|
| | def __init__(self, tanh_multiplier=None):
|
| | super(TransformerNetwork, self).__init__()
|
| | self.ConvBlock = nn.Sequential(
|
| | ConvLayer(3, 32, 9, 1),
|
| | nn.ReLU(),
|
| | ConvLayer(32, 64, 3, 2),
|
| | nn.ReLU(),
|
| | ConvLayer(64, 128, 3, 2),
|
| | nn.ReLU()
|
| | )
|
| | self.ResidualBlock = nn.Sequential(
|
| | ResidualLayer(128, 3),
|
| | ResidualLayer(128, 3),
|
| | ResidualLayer(128, 3),
|
| | ResidualLayer(128, 3),
|
| | ResidualLayer(128, 3)
|
| | )
|
| | self.DeconvBlock = nn.Sequential(
|
| | DeconvLayer(128, 64, 3, 2, 1),
|
| | nn.ReLU(),
|
| | DeconvLayer(64, 32, 3, 2, 1),
|
| | nn.ReLU(),
|
| | ConvLayer(32, 3, 9, 1, norm="None")
|
| | )
|
| | self.tanh_multiplier = tanh_multiplier
|
| |
|
| | def forward(self, x):
|
| | x = self.ConvBlock(x)
|
| | x = self.ResidualBlock(x)
|
| | x = self.DeconvBlock(x)
|
| | if isinstance(self.tanh_multiplier, int):
|
| | x = self.tanh_multiplier * F.tanh(x)
|
| | return x
|
| |
|
| | class ConvLayer(nn.Module):
|
| | def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"):
|
| | super(ConvLayer, self).__init__()
|
| | # Padding Layers
|
| | padding_size = kernel_size // 2
|
| | self.pad = nn.ReflectionPad2d(padding_size)
|
| |
|
| | # Convolution Layer
|
| | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
|
| |
|
| | # Normalization Layers
|
| | if norm == "instance":
|
| | self.norm = nn.InstanceNorm2d(out_channels, affine=True)
|
| | elif norm == "batch":
|
| | self.norm = nn.BatchNorm2d(out_channels, affine=True)
|
| | else:
|
| | self.norm = nn.Identity()
|
| |
|
| | def forward(self, x):
|
| | x = self.pad(x)
|
| | x = self.conv(x)
|
| | x = self.norm(x)
|
| | return x
|
| |
|
| | class ResidualLayer(nn.Module):
|
| | """
|
| | Deep Residual Learning for Image Recognition
|
| | https://arxiv.org/abs/1512.03385
|
| | """
|
| | def __init__(self, channels=128, kernel_size=3):
|
| | super(ResidualLayer, self).__init__()
|
| | self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1)
|
| | self.relu = nn.ReLU()
|
| | self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1)
|
| |
|
| | def forward(self, x):
|
| | identity = x # preserve residual
|
| | out = self.relu(self.conv1(x)) # 1st conv layer + activation
|
| | out = self.conv2(out) # 2nd conv layer
|
| | out = out + identity # add residual
|
| | return out
|
| |
|
| | class DeconvLayer(nn.Module):
|
| | def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"):
|
| | super(DeconvLayer, self).__init__()
|
| |
|
| | # Transposed Convolution
|
| | padding_size = kernel_size // 2
|
| | self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding_size, output_padding)
|
| |
|
| | # Normalization Layers
|
| | if norm == "instance":
|
| | self.norm = nn.InstanceNorm2d(out_channels, affine=True)
|
| | elif norm == "batch":
|
| | self.norm = nn.BatchNorm2d(out_channels, affine=True)
|
| | else:
|
| | self.norm = nn.Identity()
|
| |
|
| | def forward(self, x):
|
| | x = self.conv_transpose(x)
|
| | out = self.norm(x)
|
| | return out
|
| | ``` |
