-
Notifications
You must be signed in to change notification settings - Fork 0
/
discriminator.py
29 lines (22 loc) · 1.04 KB
/
discriminator.py
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
"""
PatchGAN Discriminator (https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py#L538)
"""
import torch.nn as nn
class Discriminator(nn.Module):
def __init__(self, args, num_filters_last=64, n_layers=3):
super(Discriminator, self).__init__()
layers = [nn.Conv2d(args.image_channels, num_filters_last, 4, 2, 1), nn.LeakyReLU(0.2)]
num_filters_mult = 1
for i in range(1, n_layers + 1):
num_filters_mult_last = num_filters_mult
num_filters_mult = min(2 ** i, 8)
layers += [
nn.Conv2d(num_filters_last * num_filters_mult_last, num_filters_last * num_filters_mult, 4,
2 if i < n_layers else 1, 1, bias=False),
nn.BatchNorm2d(num_filters_last * num_filters_mult),
nn.LeakyReLU(0.2, True)
]
layers.append(nn.Conv2d(num_filters_last * num_filters_mult, 1, 4, 1, 1))
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)