def __init__(self, ndf, nc, img_size):
super(Discriminator, self).__init__()
down_times = int(np.log2(img_size))
assert 2**down_times == img_size
max_nc = 256
self.main = nn.Sequential(
nn.utils.spectral_norm(nn.Conv2d(nc, ndf, 4, 2, 1, bias=False)),
nn.LeakyReLU(0.2, inplace=True),
# nn.InstanceNorm2d(ndf, affine=True)
)
size = int(img_size / 2)
ndf_tmp = ndf
for di in range(down_times):
if size == 2:
break
nc_out = min(2 * ndf_tmp, max_nc)
self.main.add_module(
"disc_down_" + str(size),
nn.Sequential(
nn.utils.spectral_norm(
nn.Conv2d(ndf_tmp, nc_out, 4, 2, 1, bias=False)),
nn.LeakyReLU(0.2, inplace=True), nn.BatchNorm2d(nc_out),
nn.utils.spectral_norm(
nn.Conv2d(nc_out, nc_out, 3, stride=1, padding=1)),
nn.LeakyReLU(0.2, inplace=True), nn.BatchNorm2d(nc_out)))
size = int(size / 2)
ndf_tmp = min(ndf_tmp * 2, max_nc)
out = nn.Sequential(
View(-1, 2 * 2 * ndf_tmp),
nn.utils.spectral_norm(nn.Linear(2 * 2 * ndf_tmp, ndf_tmp)),
nn.LeakyReLU(0.2, inplace=True),
nn.utils.spectral_norm(nn.Linear(ndf_tmp, 1)),
# nn.Tanh()
)
self.main.add_module("out", out)
def __init__(self):
super(VGGDiscriminator, self).__init__()
nc = 3
ndf = 64
depth = 17
self.vgg = Vgg16(depth)
self.main = nn.Sequential(
# state size. (512) x 16 x 16
nn.utils.spectral_norm(nn.Conv2d(256, ndf * 4, 4, 2, 1, bias=False)),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*4) x 8 x 8
nn.utils.spectral_norm(nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False)),
nn.BatchNorm2d(ndf * 8),
nn.LeakyReLU(0.2, inplace=True),
View(-1, ndf * 8 * 4 * 4),
nn.utils.spectral_norm(nn.Linear(ndf * 8 * 4 * 4, 10))
# state size. (ndf*8) x 4 x 4
# nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False)
)
def __init__(self, noise: Noise, image_size: int, ngf=32):
super(ResDCGenerator, self).__init__()
n_up = int(math.log2(image_size / 4))
assert 4 * (2**n_up) == image_size
nc = 3
layers = [
nn.utils.spectral_norm(
nn.Linear(noise.size(), noise.size(), bias=False)),
View(-1, noise.size(), 1, 1),
nn.utils.spectral_norm(
nn.ConvTranspose2d(noise.size(), ngf * 8, 4, 1, 0,
bias=False)),
nn.BatchNorm2d(ngf * 8),
nn.ReLU(True),
]
nc_l_next = -1
for l in range(n_up):
nc_l = max(ngf, (ngf * 8) // 2**l)
nc_l_next = max(ngf, nc_l // 2)
layers += [
Up2xResidualBlock(nc_l,
nc_l_next,
PaddingType.REFLECT,
nn.BatchNorm2d,
use_spectral_norm=True)
]
if l == 2:
layers += [SelfAttention2d(nc_l_next)]
layers += [nn.Conv2d(nc_l_next, nc, 3, 1, 1, bias=False)]
self.main = nn.Sequential(*layers)
def __init__(self,
noise: Noise,
image_size,
in_channels,
out_channels,
gen_size=32,
n_down=5,
nc_max=256):
super(UNetGenerator, self).__init__(noise)
middle_data_size = int(image_size * 2 ** (-n_down))
middle_nc = min(
int(gen_size * 2 ** n_down),
nc_max
)
assert (middle_data_size >= 4)
assert (middle_data_size % 4 == 0)
def down_block_factory(index: int) -> nn.Module:
mult = 2 ** index
in_size = min(int(gen_size * mult), nc_max)
out_size = min(2 * in_size, nc_max)
return Down2xConv2d(in_size, out_size)
def up_block_factory(index: int) -> nn.Module:
mult = 2 ** (n_down - index)
in_size = min(int(gen_size * mult), nc_max)
out_size = min(3 * in_size if index > 0 else 2 * in_size, 2 * nc_max)
return Up2xConv2d(out_size, in_size)
self.down_last_to_noise = nn.Sequential(
View(-1, middle_nc * middle_data_size**2),
spectral_norm_init(nn.Linear(middle_nc * middle_data_size**2, noise.size())),
nn.ReLU(inplace=True)
)
self.noise_up_modules = nn.Sequential(
spectral_norm_init(nn.Linear(2 * noise.size(), 2 * noise.size())),
nn.ReLU(True),
spectral_norm_init(nn.Linear(2 * noise.size(), 2 * noise.size())),
nn.ReLU(True),
View(-1, 2 * noise.size(), 1, 1),
Up4xConv2d(2 * noise.size(), middle_nc)
)
up_size = 4
while up_size < middle_data_size:
up_size *= 2
self.noise_up_modules.add_module(
"up_noize" + str(up_size),
Up2xConv2d(middle_nc, middle_nc)
)
self.unet = UNetExtra(
n_down,
in_block=nn.Sequential(
spectral_norm_init(nn.Conv2d(in_channels, gen_size, 3, stride=1, padding=1)),
nn.BatchNorm2d(gen_size),
nn.ReLU(inplace=True)
),
out_block=nn.Sequential(
spectral_norm_init(nn.Conv2d(2 * gen_size, gen_size, 3, padding=1)),
nn.BatchNorm2d(gen_size),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(gen_size, out_channels, 3, 1, 1, bias=True)
),
middle_block=self.down_last_to_noise,
middle_block_extra=self.noise_up_modules,
down_block=down_block_factory,
up_block=up_block_factory
)
def __init__(self, noise_size: int, image_size: int, ngf=32):
super(ResMeasureToImage, self).__init__()
n_up = int(math.log2(image_size / 4))
assert 4 * (2**n_up) == image_size
nc = 3
layers = [
spectral_norm_init(nn.Linear(noise_size,
noise_size // 2,
bias=False),
n_power_iterations=10),
nn.LeakyReLU(0.2, inplace=True), #nn.ReLU(inplace=True),
spectral_norm_init(nn.Linear(noise_size // 2,
noise_size // 2,
bias=False),
n_power_iterations=10),
nn.LeakyReLU(0.2, inplace=True),
View(-1, noise_size // 2, 1, 1),
spectral_norm_init(nn.ConvTranspose2d(noise_size // 2,
ngf * 4,
4,
1,
0,
bias=False),
n_power_iterations=10),
nn.InstanceNorm2d(ngf * 4), # nn.InstanceNorm2d(ngf * 8),
nn.LeakyReLU(0.2, inplace=True),
]
nc_l_next = -1
for l in range(n_up):
nc_l = max(ngf // 2, (ngf * 4) // 2**l)
nc_l_next = max(ngf // 2, nc_l // 2)
layers += [
Up2xResidualBlock(nc_l,
nc_l_next,
PaddingType.REFLECT,
nn.InstanceNorm2d,
use_spectral_norm=True,
activation=nn.LeakyReLU(
0.2, inplace=True)), #nn.InstanceNorm2d
PooledResidualBlock(nc_l_next,
nc_l_next,
nc_l_next,
nn.Identity(),
PaddingType.REFLECT,
nn.InstanceNorm2d,
use_spectral_norm=True,
activation=nn.LeakyReLU(0.2, inplace=True))
]
if l == 2:
layers += [
# nn.Dropout2d(p=0.5),
SelfAttention2d(nc_l_next)
]
layers += [nn.Conv2d(nc_l_next, nc, 3, 1, 1, bias=False)]
self.main = nn.Sequential(*layers)