Spaces:
Running
Running
File size: 3,094 Bytes
62a43b6 |
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 |
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, downsample: bool = True, use_act: bool = True,
use_dropout: bool = False, **kwargs):
super(ConvBlock, self).__init__()
self.conv_block = nn.Sequential(
nn.Conv2d(in_channels=in_channels, out_channels=out_channels, padding_mode="reflect", **kwargs)
if downsample
else nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, **kwargs),
nn.InstanceNorm2d(num_features=out_channels),
nn.ReLU(inplace=True) if use_act else nn.Identity()
)
if use_dropout:
self.conv_block = nn.Sequential(self.conv_block, nn.Dropout(p=0.5))
def forward(self, x):
return self.conv_block(x)
class ResidualBlock(nn.Module):
def __init__(self, features: int):
super(ResidualBlock, self).__init__()
self.residual_block = nn.Sequential(
ConvBlock(in_channels=features, out_channels=features, kernel_size=3, padding=1),
ConvBlock(in_channels=features, out_channels=features, kernel_size=3, padding=1, use_act=False),
)
def forward(self, x):
return x + self.residual_block(x)
class CycleGenerator(nn.Module):
def __init__(self, img_channels: int = 3, latent_dim: int = 64, num_residuals: int = 9):
super(CycleGenerator, self).__init__()
self.base = nn.Sequential(
nn.Conv2d(in_channels=img_channels, out_channels=latent_dim, kernel_size=7, stride=1, padding=3,
padding_mode="reflect"),
nn.ReLU(inplace=True)
)
self.down_blocks = nn.ModuleList(
[
ConvBlock(in_channels=latent_dim, out_channels=latent_dim * 2, kernel_size=3, stride=2, padding=1),
ConvBlock(in_channels=latent_dim * 2, out_channels=latent_dim * 4, kernel_size=3, stride=2, padding=1),
]
)
self.residual_blocks = nn.Sequential(
*[ResidualBlock(features=latent_dim * 4) for _ in range(num_residuals)]
)
self.up_blocks = nn.ModuleList(
[
ConvBlock(in_channels=latent_dim * 4, out_channels=latent_dim * 2, kernel_size=3, stride=2, padding=1,
output_padding=1,
downsample=False),
ConvBlock(in_channels=latent_dim * 2, out_channels=latent_dim, kernel_size=3, stride=2, padding=1,
output_padding=1,
downsample=False),
]
)
self.head = nn.Conv2d(in_channels=latent_dim, out_channels=img_channels, kernel_size=7, stride=1, padding=3,
padding_mode="reflect")
def forward(self, x):
x = self.base(x)
for layer in self.down_blocks:
x = layer(x)
x = self.residual_blocks(x)
for layer in self.up_blocks:
x = layer(x)
x = self.head(x)
return torch.tanh(x)
|