122 lines
4.6 KiB
Python
122 lines
4.6 KiB
Python
import functools
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
import torch
|
|
|
|
|
|
def weights_init(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find('Conv') != -1:
|
|
nn.init.normal_(m.weight.data, 0.0, 0.02)
|
|
nn.init.constant_(m.bias.data, 0)
|
|
elif classname.find('BatchNorm') != -1:
|
|
nn.init.normal_(m.weight.data, 1.0, 0.02)
|
|
nn.init.constant_(m.bias.data, 0)
|
|
|
|
|
|
class NLayerDiscriminator(nn.Module):
|
|
"""Defines a PatchGAN discriminator as in Pix2Pix
|
|
--> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
|
|
"""
|
|
def __init__(self, input_nc=3, ndf=64, n_layers=4):
|
|
"""Construct a PatchGAN discriminator
|
|
Parameters:
|
|
input_nc (int) -- the number of channels in input images
|
|
ndf (int) -- the number of filters in the last conv layer
|
|
n_layers (int) -- the number of conv layers in the discriminator
|
|
norm_layer -- normalization layer
|
|
"""
|
|
super(NLayerDiscriminator, self).__init__()
|
|
|
|
# norm_layer = nn.BatchNorm2d
|
|
norm_layer = nn.InstanceNorm2d
|
|
|
|
if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
|
|
use_bias = norm_layer.func != nn.BatchNorm2d
|
|
else:
|
|
use_bias = norm_layer != nn.BatchNorm2d
|
|
|
|
kw = 4
|
|
padw = 1
|
|
sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
|
|
nf_mult = 1
|
|
nf_mult_prev = 1
|
|
for n in range(1, n_layers): # gradually increase the number of filters
|
|
nf_mult_prev = nf_mult
|
|
nf_mult = min(2 ** n, 8)
|
|
sequence += [
|
|
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
|
|
norm_layer(ndf * nf_mult),
|
|
nn.LeakyReLU(0.2, True)
|
|
]
|
|
|
|
nf_mult_prev = nf_mult
|
|
nf_mult = min(2 ** n_layers, 8)
|
|
sequence += [
|
|
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
|
|
norm_layer(ndf * nf_mult),
|
|
nn.LeakyReLU(0.2, True)
|
|
]
|
|
|
|
sequence += [
|
|
nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
|
|
self.main = nn.Sequential(*sequence)
|
|
|
|
def forward(self, input):
|
|
"""Standard forward."""
|
|
return self.main(input)
|
|
|
|
|
|
class NLayerDiscriminator3D(nn.Module):
|
|
"""Defines a 3D PatchGAN discriminator as in Pix2Pix but for 3D inputs."""
|
|
def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
|
|
"""
|
|
Construct a 3D PatchGAN discriminator
|
|
|
|
Parameters:
|
|
input_nc (int) -- the number of channels in input volumes
|
|
ndf (int) -- the number of filters in the last conv layer
|
|
n_layers (int) -- the number of conv layers in the discriminator
|
|
use_actnorm (bool) -- flag to use actnorm instead of batchnorm
|
|
"""
|
|
super(NLayerDiscriminator3D, self).__init__()
|
|
# if not use_actnorm:
|
|
# norm_layer = nn.BatchNorm3d
|
|
# else:
|
|
# raise NotImplementedError("Not implemented.")
|
|
|
|
norm_layer = nn.InstanceNorm3d
|
|
|
|
if type(norm_layer) == functools.partial:
|
|
use_bias = norm_layer.func != nn.BatchNorm3d
|
|
else:
|
|
use_bias = norm_layer != nn.BatchNorm3d
|
|
|
|
kw = 4
|
|
padw = 1
|
|
sequence = [nn.Conv3d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
|
|
nf_mult = 1
|
|
nf_mult_prev = 1
|
|
for n in range(1, n_layers): # gradually increase the number of filters
|
|
nf_mult_prev = nf_mult
|
|
nf_mult = min(2 ** n, 8)
|
|
sequence += [
|
|
nn.Conv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=(1,2,2), padding=padw, bias=use_bias),
|
|
norm_layer(ndf * nf_mult),
|
|
nn.LeakyReLU(0.2, True)
|
|
]
|
|
|
|
nf_mult_prev = nf_mult
|
|
nf_mult = min(2 ** n_layers, 8)
|
|
sequence += [
|
|
nn.Conv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=1, padding=padw, bias=use_bias),
|
|
norm_layer(ndf * nf_mult),
|
|
nn.LeakyReLU(0.2, True)
|
|
]
|
|
|
|
sequence += [nn.Conv3d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
|
|
self.main = nn.Sequential(*sequence)
|
|
|
|
def forward(self, input):
|
|
"""Standard forward."""
|
|
return self.main(input) |