Spaces:
Runtime error
Runtime error
File size: 5,460 Bytes
0587b57 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """Model definitions for the CycleGAN-style architecture."""
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
class ResidualBlock(nn.Module):
"""Simple residual block with two conv layers."""
def __init__(self, in_features: int) -> None:
super().__init__()
conv_block = [
nn.ReflectionPad2d(1), # (B, C, H+2, W+2)
nn.Conv2d(in_features, in_features, 3), # (B, C, H, W)
nn.BatchNorm2d(in_features), # (B, C, H, W)
nn.ReLU(), # (B, C, H, W)
nn.ReflectionPad2d(1), # (B, C, H+2, W+2)
nn.Conv2d(in_features, in_features, 3), # (B, C, H, W)
nn.BatchNorm2d(in_features),
] # (B, C, H, W)
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x: Tensor) -> Tensor:
"""Apply the residual block."""
return x + self.conv_block(x)
class Generator(nn.Module):
"""U-Net style generator used for domain translation."""
def __init__(self, ngf: int, n_residual_blocks: int = 9) -> None:
super().__init__()
# Initial convlution block
model = [
nn.ReflectionPad2d(
3
), # (B, 3, H+6, W+6), applies 2D "reflection" padding of 3 pixels on all four sides of image
nn.Conv2d(
3, ngf, 7
), # (B, ngf, H, W), 3 in_channels, ngf out_channels, kernel size 7 (keeps same image size)
nn.BatchNorm2d(
ngf
), # (B, ngf, H, W), normalized for each ngf across all B, H, W
nn.ReLU(),
] # (B, ngf, H, W)
# Downsampling
in_features = ngf # number of generator filters
out_features = in_features * 2
for _ in range(2):
model += [
nn.Conv2d(
in_features, out_features, 3, stride=2, padding=1
), # (B, in_features*2, H//2, W//2), doubles number of channels and reduces H, W by half
nn.BatchNorm2d(out_features), # (B, in_features*2, H//2, W//2)
nn.ReLU(),
] # (B, in_features*2, H//2, W//2)
in_features = out_features
out_features = in_features * 2
# Residual blocks
for _ in range(n_residual_blocks):
model += [
ResidualBlock(in_features)
] # (B, in_features, H, W), returns same size as input
# Upsampling
out_features = in_features // 2
for _ in range(2):
model += [
nn.ConvTranspose2d(
in_features, out_features, 3, stride=2, padding=1, output_padding=1
), # (B, in_features//2, H*2, W*2), upsamples to twice the H, W with half the channels
nn.BatchNorm2d(out_features), # (B, in_features//2, H*2, W*2)
nn.ReLU(),
] # (B, in_features//2, H*2, W*2)
in_features = out_features
out_features = in_features // 2
# Output layer
model += [
nn.ReflectionPad2d(3), # (B, in_features, H+6, W+6)
nn.Conv2d(ngf, 3, 7), # (B, 3, H, W)
nn.Tanh(),
] # (B, 3, H, W), passed tanh activation
self.model = nn.Sequential(*model)
def forward(self, x: Tensor) -> Tensor:
"""Generate an image from ``x``."""
return self.model(x)
class Discriminator(nn.Module):
"""PatchGAN discriminator."""
def __init__(self, ndf: int) -> None:
super().__init__()
model = [
nn.Conv2d(
3, ndf, 4, stride=2, padding=1
), # (B, ndf, H//2, W//2), channel from 3 -> ndf
nn.LeakyReLU(0.2, inplace=True),
] # (B, ndf, H//2, W//2)
model += [
nn.Conv2d(ndf, ndf * 2, 4, stride=2, padding=1), # (B, ndf * 2, H//4, W//4)
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
]
model += [
nn.Conv2d(
ndf * 2, ndf * 4, 4, stride=2, padding=1
), # (B, ndf * 4, H//8, W//8)
nn.InstanceNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
]
model += [
nn.Conv2d(ndf * 4, ndf * 8, 4, padding=1), # (B, ndf * 8, H//8-1, W//8-1)
nn.InstanceNorm2d(ndf * 8),
nn.LeakyReLU(0.2, inplace=True),
]
# FCN classification layer
model += [nn.Conv2d(ndf * 8, 1, 4, padding=1)] # (B, 1, H//8-2, W//8-2)
self.model = nn.Sequential(*model)
def forward(self, x: Tensor) -> Tensor:
"""Return discriminator logits for input ``x``."""
# x: (B, 3, H, W)
x = self.model(x) # (B, 1, H//8-2, W//8-2)
# Average pooling and flatten
return F.avg_pool2d(x, x.size()[2:]).view(
x.size()[0], -1
) # global average -> (B, 1, 1, 1) -> flatten to (B, 1)
# Initialize and return the generators and discriminators used for training
def initialize_models(
ngf: int = 32,
ndf: int = 32,
n_blocks: int = 9,
) -> tuple[Generator, Generator, Discriminator, Discriminator]:
"""Instantiate generators and discriminators with default sizes."""
# initialize the generators and discriminators
G = Generator(ngf, n_blocks)
F = Generator(ngf, n_blocks)
DX = Discriminator(ndf)
DY = Discriminator(ndf)
return G, F, DX, DY
|