File size: 5,236 Bytes
e5b70eb |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
class DiscriminatorVGG(nn.Module):
def __init__(self, in_ch=3, image_size=128, d=64):
super(DiscriminatorVGG, self).__init__()
self.feature_map_size = image_size // 32
self.d = d
self.features = nn.Sequential(
nn.Conv2d(in_ch, d, kernel_size=3, stride=1, padding=1), # input is 3 x 128 x 128
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d, d, kernel_size=3, stride=2, padding=1, bias=False), # state size. 64 x 64 x 64
nn.BatchNorm2d(d),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d, d*2, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(d*2),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*2, d*2, kernel_size=3, stride=2, padding=1, bias=False), # state size. 128 x 32 x 32
nn.BatchNorm2d(d*2),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*2, d*4, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(d*4),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*4, d*4, kernel_size=3, stride=2, padding=1, bias=False), # state size. 256 x 16 x 16
nn.BatchNorm2d(d*4),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*4, d*8, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(d*8),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*8, d*8, kernel_size=3, stride=2, padding=1, bias=False), # state size. 512 x 8 x 8
nn.BatchNorm2d(d*8),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*8, d*8, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(d*8),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(d*8, d*8, kernel_size=3, stride=2, padding=1, bias=False), # state size. 512 x 4 x 4
nn.BatchNorm2d(d*8),
nn.LeakyReLU(negative_slope=0.2, inplace=True)
)
self.classifier = nn.Sequential(
nn.Linear((self.d*8) * self.feature_map_size * self.feature_map_size, 100),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Linear(100, 1)
)
def forward(self, x):
out = self.features(x)
out = torch.flatten(out, 1)
out = self.classifier(out)
return out
class UNetDiscriminator(nn.Module):
def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
super(UNetDiscriminator, self).__init__()
self.skip_connection = skip_connection
norm = spectral_norm
self.num_in_ch = num_in_ch
self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
self.conv1 = norm(nn.Conv2d(num_feat, num_feat*2, kernel_size=4, stride=2, padding=1, bias=False))
self.conv2 = norm(nn.Conv2d(num_feat*2, num_feat*4, kernel_size=4, stride=2, padding=1, bias=False))
self.conv3 = norm(nn.Conv2d(num_feat*4, num_feat*8, kernel_size=4, stride=2, padding=1, bias=False))
# upsample
self.conv4 = norm(nn.Conv2d(num_feat*8, num_feat*4, kernel_size=3, stride=1, padding=1, bias=False))
self.conv5 = norm(nn.Conv2d(num_feat*4, num_feat*2, kernel_size=3, stride=1, padding=1, bias=False))
self.conv6 = norm(nn.Conv2d(num_feat*2, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
# extra
self.conv7 = norm(nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
self.conv8 = norm(nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
self.conv9 = nn.Conv2d(num_feat, 1, kernel_size=3, stride=1, padding=1)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
def forward(self, x):
x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
# upsample
x3 = F.interpolate(x3, scale_factor=2, mode='bilinear')
x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
if self.skip_connection:
x4 = x4 + x2
x4 = F.interpolate(x4, scale_factor=2, mode='bilinear')
x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
if self.skip_connection:
x5 = x5 + x1
x5 = F.interpolate(x5, scale_factor=2, mode='bilinear')
x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
if self.skip_connection:
x6 = x6 + x0
# extra
out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
out = self.conv9(out)
out = self.avg_pool(out)
out = torch.flatten(out, 1)
return out
|