| import torch |
| import torch.nn as nn |
| |
| class ResidualBlock(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.block = nn.Sequential( |
| nn.ReflectionPad2d(1), |
| nn.Conv2d(channels, channels, 3), |
| nn.InstanceNorm2d(channels), |
| nn.ReLU(inplace=True), |
| nn.ReflectionPad2d(1), |
| nn.Conv2d(channels, channels, 3), |
| nn.InstanceNorm2d(channels), |
| ) |
| def forward(self, x): |
| return x + self.block(x) |
|
|
| class SelfAttention(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.query = nn.Conv2d(channels, channels // 8, 1) |
| self.key = nn.Conv2d(channels, channels // 8, 1) |
| self.value = nn.Conv2d(channels, channels, 1) |
| self.gamma = nn.Parameter(torch.zeros(1)) |
|
|
| def forward(self, x): |
| B, C, H, W = x.shape |
| q = self.query(x).flatten(2) |
| k = self.key(x).flatten(2) |
| v = self.value(x).flatten(2) |
| attn = torch.softmax(torch.bmm(q.transpose(1,2), k), dim=-1) |
| out = torch.bmm(v, attn.transpose(1,2)).view(B, C, H, W) |
| return x + self.gamma * out |
|
|
| class ResNetGenerator(nn.Module): |
| def __init__(self, in_channels=3, out_channels=3, n_filters=64, n_res_blocks=9): |
| super().__init__() |
| model = [ |
| nn.ReflectionPad2d(3), |
| nn.Conv2d(in_channels, n_filters, 7), |
| nn.InstanceNorm2d(n_filters), |
| nn.ReLU(inplace=True), |
|
|
| nn.Conv2d(n_filters, n_filters*2, 3, stride=2, padding=1), |
| nn.InstanceNorm2d(n_filters*2), |
| nn.ReLU(inplace=True), |
|
|
| nn.Conv2d(n_filters*2, n_filters*4, 3, stride=2, padding=1), |
| nn.InstanceNorm2d(n_filters*4), |
| nn.ReLU(inplace=True), |
| ] |
|
|
| for _ in range(n_res_blocks): |
| model.append(ResidualBlock(n_filters*4)) |
|
|
| model.append(SelfAttention(n_filters*4)) |
|
|
| model += [ |
| nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), |
| nn.Conv2d(n_filters*4, n_filters*2, 3, padding=1), |
| nn.InstanceNorm2d(n_filters*2), |
| nn.ReLU(inplace=True), |
|
|
| nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), |
| nn.Conv2d(n_filters*2, n_filters, 3, padding=1), |
| nn.InstanceNorm2d(n_filters), |
| nn.ReLU(inplace=True), |
|
|
| nn.ReflectionPad2d(3), |
| nn.Conv2d(n_filters, out_channels, 7), |
| nn.Tanh() |
| ] |
|
|
| self.model = nn.Sequential(*model) |
|
|
| def forward(self, x): |
| return self.model(x) |
| |
| @torch.no_grad() |
| def load_generator(path, device="cpu"): |
| gen = ResNetGenerator() |
| state_dict = torch.load(path, map_location="cpu") |
| state_dict = {k: v.float() for k, v in state_dict.items()} |
| gen.load_state_dict(state_dict) |
| gen.to(device).eval() |
| return gen |
|
|
