Spaces:
Sleeping
Sleeping
File size: 2,259 Bytes
5a14c00 442d9ae | 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 | import torch
import torch.nn as nn
class ResidualBlock(nn.Module):
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(channels, channels, 3, padding=1),
nn.InstanceNorm2d(channels),
nn.ReLU(),
nn.Conv2d(channels, channels, 3, padding=1),
nn.InstanceNorm2d(channels)
)
def forward(self, x):
return x + self.conv(x)
class UpsampleConvLayer(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride):
super().__init__()
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
# Reflection padding reduces "border" artifacts
self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)
def forward(self, x):
x = self.upsample(x)
x = self.reflection_pad(x)
return self.conv(x)
class StyleNet(nn.Module):
def __init__(self):
super(StyleNet, self).__init__()
# Encoder
self.encoder = nn.Sequential(
nn.ReflectionPad2d(4), # Pad first...
nn.Conv2d(3, 32, 9, padding=0), # ...then Conv (padding=0)
nn.InstanceNorm2d(32), nn.ReLU(),
nn.ReflectionPad2d(1),
nn.Conv2d(32, 64, 3, stride=2, padding=0),
nn.InstanceNorm2d(64), nn.ReLU(),
nn.ReflectionPad2d(1),
nn.Conv2d(64, 128, 3, stride=2, padding=0),
nn.InstanceNorm2d(128), nn.ReLU()
)
# Bottleneck
self.bottleneck = nn.Sequential(*[ResidualBlock(128) for _ in range(3)])
# Decoder
self.decoder = nn.Sequential(
UpsampleConvLayer(128, 64, kernel_size=3, stride=1),
nn.InstanceNorm2d(64), nn.ReLU(),
UpsampleConvLayer(64, 32, kernel_size=3, stride=1),
nn.InstanceNorm2d(32), nn.ReLU(),
nn.Conv2d(32, 3, 9, padding=4),
nn.Sigmoid()
)
def forward(self, x):
features = self.encoder(x)
features = self.bottleneck(features)
return self.decoder(features) |