File size: 3,188 Bytes
10842a1 | 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 64 65 66 67 68 69 70 71 | import torch
import torch.nn as nn
class ResidualBlock(nn.Module):
def __init__(self, in_features):
super(ResidualBlock, self).__init__()
self.block = nn.Sequential(
nn.ReflectionPad2d(1), nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features), nn.ReLU(inplace=True),
nn.ReflectionPad2d(1), nn.Conv2d(in_features, in_features, 3),
nn.InstanceNorm2d(in_features)
)
def forward(self, x):
return x + self.block(x)
class ResNetGenerator(nn.Module):
def __init__(self, input_channels=3, output_channels=3, num_residual_blocks=9):
super(ResNetGenerator, self).__init__()
out_features = 64
model =[nn.ReflectionPad2d(3), nn.Conv2d(input_channels, out_features, 7), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True)]
in_features = out_features
for _ in range(2):
out_features *= 2
model +=[nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True)]
in_features = out_features
for _ in range(num_residual_blocks):
model += [ResidualBlock(in_features)]
for _ in range(2):
out_features //= 2
model +=[nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True)]
in_features = out_features
model +=[nn.ReflectionPad2d(3), nn.Conv2d(out_features, output_channels, 7), nn.Tanh()]
self.model = nn.Sequential(*model)
def forward(self, x):
return self.model(x)
class PatchGANDiscriminator(nn.Module):
def __init__(self, input_channels=3):
super(PatchGANDiscriminator, self).__init__()
def discriminator_block(in_filters, out_filters, stride=2, normalize=True):
layers =[nn.Conv2d(in_filters, out_filters, 4, stride=stride, padding=1)]
if normalize: layers.append(nn.InstanceNorm2d(out_filters))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
self.model = nn.Sequential(
*discriminator_block(input_channels, 64, normalize=False),
*discriminator_block(64, 128),
*discriminator_block(128, 256),
*discriminator_block(256, 512, stride=1),
nn.Conv2d(512, 1, 4, padding=1)
)
def forward(self, x):
return self.model(x)
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0.0, 0.02)
if hasattr(m, 'bias') and m.bias is not None:
nn.init.constant_(m.bias.data, 0.0)
class CycleGAN(nn.Module):
def __init__(self):
super(CycleGAN, self).__init__()
self.G_A2B = ResNetGenerator(num_residual_blocks=9)
self.G_B2A = ResNetGenerator(num_residual_blocks=9)
self.D_A = PatchGANDiscriminator()
self.D_B = PatchGANDiscriminator()
self.apply(weights_init_normal)
|