def __init__(self, dim, use_bias):
super(ResnetBlock, self).__init__()
conv_block = []
conv_block += [
nn.ReflectionPad2d(1),
nn.Conv2d(dim,
dim,
kernel_size=3,
stride=1,
padding=0,
bias=use_bias),
nn.InstanceNorm2d(dim, affine=True),
nn.ReLU(True)
]
conv_block += [
nn.ReflectionPad2d(1),
nn.Conv2d(dim,
dim,
kernel_size=3,
stride=1,
padding=0,
bias=use_bias),
nn.InstanceNorm2d(dim, affine=True)
]
self.conv_block = nn.Sequential(*conv_block)
def __init__(self, num_classes=1000):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def __init__(self, input_nc, ndf=64, n_layers=5):
super(Discriminator, self).__init__()
model = [nn.ReflectionPad2d(1),
nn.utils.spectral_norm(nn.Conv2d(3, ndf, 4, 2, 0, bias=True)),
nn.LeakyReLU(0.2, True)]
for i in range(1, n_layers - 2):
mult = 2 ** (i - 1)
model += [nn.ReflectionPad2d(1),
nn.utils.spectral_norm(nn.Conv2d(ndf * mult, ndf * mult * 2, 4, 2, 0, bias=True)),
nn.LeakyReLU(0.2, True)]
mult = 2 ** (n_layers - 2 - 1)
model += [nn.ReflectionPad2d(1),
nn.utils.spectral_norm(nn.Conv2d(ndf * mult, ndf * mult * 2, 4, 1, 0, bias=True)),
nn.LeakyReLU(0.2, True)]
# Class Activation Map
mult = 2 ** (n_layers - 2)
self.gap_fc = nn.utils.spectral_norm(nn.Linear(ndf * mult, 1, bias=False))
self.gmp_fc = nn.utils.spectral_norm(nn.Linear(ndf * mult, 1, bias=False))
self.conv1x1 = nn.Conv2d(ndf * mult * 2, ndf * mult, 1, 1, bias=True)
self.leaky_relu = nn.LeakyReLU(0.2, True)
self.pad = nn.ReflectionPad2d(1)
self.conv = nn.utils.spectral_norm(nn.Conv2d(ndf * mult, 1, 4, 1, 0, bias=False))
self.model = nn.Sequential(*model)
def _build_weights(self, dim_in, dim_out, style_dim=64):
self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
self.norm1 = AdaIN(style_dim, dim_in)
self.norm2 = AdaIN(style_dim, dim_out)
if self.learned_sc:
self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
def __init__(self,
num_modules=1,
end_relu=False,
num_landmarks=98,
fname_pretrained=None):
super(FAN, self).__init__()
self.num_modules = num_modules
self.end_relu = end_relu
# Base part
self.conv1 = CoordConvTh(256,
256,
True,
False,
in_channels=3,
out_channels=64,
kernel_size=7,
stride=2,
padding=3)
self.bn1 = nn.BatchNorm2d(64)
self.conv2 = ConvBlock(64, 128)
self.conv3 = ConvBlock(128, 128)
self.conv4 = ConvBlock(128, 256)
# Stacking part
self.add_module('m0', HourGlass(1, 4, 256, first_one=True))
self.add_module('top_m_0', ConvBlock(256, 256))
self.add_module('conv_last0', nn.Conv2d(256, 256, 1, 1, 0))
self.add_module('bn_end0', nn.BatchNorm2d(256))
self.add_module('l0', nn.Conv2d(256, num_landmarks + 1, 1, 1, 0))
if fname_pretrained is not None:
self.load_pretrained_weights(fname_pretrained)
def _build_weights(self, dim_in, dim_out):
self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
if self.normalize:
self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
if self.learned_sc:
self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
def __init__(self, dim):
super(ResNetAdaILNBlock, self).__init__()
self.pad1 = nn.ReflectionPad2d(1)
self.conv1 = nn.Conv2d(dim, dim, 3, 1, 0, bias=False)
self.norm1 = AdaILN(dim)
self.relu1 = nn.ReLU(True)
self.pad2 = nn.ReflectionPad2d(1)
self.conv2 = nn.Conv2d(dim, dim, 3, 1, 0, bias=False)
self.norm2 = AdaILN(dim)
def __init__(self, dim):
super(ResNetBlock, self).__init__()
conv_block = []
conv_block += [nn.ReflectionPad2d(1),
nn.Conv2d(dim, dim, 3, 1, 0, bias=False),
nn.InstanceNorm2d(dim,affine=True),
nn.ReLU(True)]
conv_block += [nn.ReflectionPad2d(1),
nn.Conv2d(dim, dim, 3, 1, 0, bias=False),
nn.InstanceNorm2d(dim,affine=True)]
self.conv_block = nn.Sequential(*conv_block)
def __init__(self, img_size=256, num_domains=2, max_conv_dim=512):
super().__init__()
dim_in = 2**14 // img_size
blocks = []
blocks += [nn.Conv2d(3, dim_in, 3, 1, 1)]
repeat_num = int(np.log2(img_size)) - 2
for _ in range(repeat_num):
dim_out = min(dim_in * 2, max_conv_dim)
blocks += [ResBlk(dim_in, dim_out, downsample=True)]
dim_in = dim_out
blocks += [nn.LeakyReLU(0.2)]
blocks += [nn.Conv2d(dim_out, dim_out, 4, 1, 0)]
blocks += [nn.LeakyReLU(0.2)]
blocks += [nn.Conv2d(dim_out, num_domains, 1, 1, 0)]
self.main = nn.Sequential(*blocks)
def __init__(self, dim, use_bias):
super(ResnetAdaILNBlock, self).__init__()
self.pad1 = nn.ReflectionPad2d(1)
self.conv1 = nn.Conv2d(dim,
dim,
kernel_size=3,
stride=1,
padding=0,
bias=use_bias)
self.norm1 = adaILN(dim)
self.relu1 = nn.ReLU(True)
self.pad2 = nn.ReflectionPad2d(1)
self.conv2 = nn.Conv2d(dim,
dim,
kernel_size=3,
stride=1,
padding=0,
bias=use_bias)
self.norm2 = adaILN(dim)
def __init__(self, img_size=256, style_dim=64, max_conv_dim=512, w_hpf=1):
super().__init__()
dim_in = 2**14 // img_size
self.img_size = img_size
self.from_rgb = nn.Conv2d(3, dim_in, 3, 1, 1)
self.encode = nn.ModuleList()
self.decode = nn.ModuleList()
self.to_rgb = nn.Sequential(nn.InstanceNorm2d(dim_in, affine=True),
nn.LeakyReLU(0.2),
nn.Conv2d(dim_in, 3, 1, 1, 0))
# down/up-sampling blocks
repeat_num = int(np.log2(img_size)) - 4
if w_hpf > 0:
repeat_num += 1
for _ in range(repeat_num):
dim_out = min(dim_in * 2, max_conv_dim)
self.encode.append(
ResBlk(dim_in, dim_out, normalize=True, downsample=True))
self.decode.insert(0,
AdainResBlk(dim_out,
dim_in,
style_dim,
w_hpf=w_hpf,
upsample=True)) # stack-like
dim_in = dim_out
# bottleneck blocks
for _ in range(2):
self.encode.append(ResBlk(dim_out, dim_out, normalize=True))
self.decode.insert(
0, AdainResBlk(dim_out, dim_out, style_dim, w_hpf=w_hpf))
if w_hpf > 0:
device = porch.device(
'cuda' if porch.cuda.is_available() else 'cpu')
self.hpf = HighPass(w_hpf, device)
def __init__(self,
height,
width,
with_r,
with_boundary,
in_channels,
first_one=False,
*args,
**kwargs):
super(CoordConvTh, self).__init__()
self.addcoords = AddCoordsTh(height, width, with_r, with_boundary)
in_channels += 2
if with_r:
in_channels += 1
if with_boundary and not first_one:
in_channels += 2
self.conv = nn.Conv2d(in_channels=in_channels, *args, **kwargs)
def __init__(self, in_planes, out_planes):
super(ConvBlock, self).__init__()
self.bn1 = nn.BatchNorm2d(in_planes)
conv3x3 = partial(nn.Conv2d,
kernel_size=3,
stride=1,
padding=1,
bias=False,
dilation=1)
self.conv1 = conv3x3(in_planes, int(out_planes / 2))
self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4))
self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4))
self.downsample = None
if in_planes != out_planes:
self.downsample = nn.Sequential(
nn.BatchNorm2d(in_planes), nn.ReLU(True),
nn.Conv2d(in_planes, out_planes, 1, 1, bias=False))
def __init__(self, input_nc, output_nc, ngf=64, n_blocks=6, img_size=256, light=False):
super(ResnetGenerator, self).__init__()
self.n_res=n_blocks
self.light= light
down_layer = [
nn.ReflectionPad2d(3),
nn.Conv2d(3, ngf, 7, 1, 0, bias=False),
nn.InstanceNorm2d(ngf,affine=True),
nn.ReLU(inplace=True),
# Down-Sampling
nn.ReflectionPad2d(1),
nn.Conv2d(ngf, ngf*2, 3, 2, 0, bias=False),
nn.InstanceNorm2d(ngf*2,affine=True),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(1),
nn.Conv2d(ngf*2, ngf*4, 3, 2, 0, bias=False),
nn.InstanceNorm2d(ngf*4,affine=True),
nn.ReLU(inplace=True),
# Down-Sampling Bottleneck
ResNetBlock(ngf*4),
ResNetBlock(ngf*4),
ResNetBlock(ngf*4),
ResNetBlock(ngf*4),
]
# Class Activation Map
self.gap_fc = nn.Linear(ngf*4, 1, bias=False)
self.gmp_fc = nn.Linear(ngf*4, 1, bias=False)
self.conv1x1 = nn.Conv2d(ngf*8, ngf*4, 1, 1, bias=True)
self.relu = nn.ReLU(inplace=True)
# # Gamma, Beta block
# fc = [
# nn.Linear(image_size * image_size * 16, 256, bias=False),
# nn.ReLU(inplace=True),
# nn.Linear(256, 256, bias=False),
# nn.ReLU(inplace=True)
# ]
# Gamma, Beta block
if self.light:
fc = [nn.Linear(ngf*4, ngf*4, bias=False),
nn.ReLU(True),
nn.Linear(ngf*4, ngf*4, bias=False),
nn.ReLU(True)]
else:
fc = [nn.Linear(img_size * img_size * ngf//4, ngf*4, bias=False),
nn.ReLU(True),
nn.Linear(ngf*4, ngf*4, bias=False),
nn.ReLU(True)]
self.gamma = nn.Linear(ngf*4, ngf*4, bias=False)
self.beta = nn.Linear(ngf*4, ngf*4, bias=False)
# Up-Sampling Bottleneck
for i in range(self.n_res):
setattr(self, "ResNetAdaILNBlock_" + str(i + 1), ResNetAdaILNBlock(ngf*4))
up_layer = [
nn.Upsample(scale_factor=2, mode="nearest"),
nn.ReflectionPad2d(1),
nn.Conv2d(ngf*4, ngf*2, 3, 1, 0, bias=False),
ILN(ngf*2),
nn.ReLU(inplace=True),
nn.Upsample(scale_factor=2, mode="nearest"),
nn.ReflectionPad2d(1),
nn.Conv2d(ngf*2, ngf, 3, 1, 0, bias=False),
ILN(ngf),
nn.ReLU(inplace=True),
nn.ReflectionPad2d(3),
nn.Conv2d(ngf, 3, 7, 1, 0, bias=False),
nn.Tanh()
]
self.down_layer = nn.Sequential(*down_layer)
self.fc = nn.Sequential(*fc)
self.up_layer = nn.Sequential(*up_layer)
def __init__(self, in_channels, out_channels, **kwargs):
super(BasicConv2d, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
def __init__(self, in_channels, out_channels=1):
super().__init__()
self.main = nn.Sequential(
nn.Dropout(0.5),
nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False))
def __init__(self,
input_nc,
output_nc,
ngf=64,
n_blocks=6,
img_size=256,
light=False):
assert (n_blocks >= 0)
super(ResnetGenerator, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
self.n_blocks = n_blocks
self.img_size = img_size
self.light = light
DownBlock = []
DownBlock += [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc,
ngf,
kernel_size=7,
stride=1,
padding=0,
bias=False),
nn.InstanceNorm2d(ngf, affine=True),
nn.ReLU(True)
]
# Down-Sampling
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
DownBlock += [
nn.ReflectionPad2d(1),
nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=0,
bias=False),
nn.InstanceNorm2d(ngf * mult * 2, affine=True),
nn.ReLU(True)
]
# Down-Sampling Bottleneck
mult = 2**n_downsampling
for i in range(n_blocks):
DownBlock += [ResnetBlock(ngf * mult, use_bias=False)]
# Class Activation Map
self.gap_fc = nn.Linear(ngf * mult, 1, bias=False)
self.gmp_fc = nn.Linear(ngf * mult, 1, bias=False)
self.conv1x1 = nn.Conv2d(ngf * mult * 2,
ngf * mult,
kernel_size=1,
stride=1,
bias=True)
self.relu = nn.ReLU(True)
# Gamma, Beta block
if self.light:
FC = [
nn.Linear(ngf * mult, ngf * mult, bias=False),
nn.ReLU(True),
nn.Linear(ngf * mult, ngf * mult, bias=False),
nn.ReLU(True)
]
else:
FC = [
nn.Linear(img_size // mult * img_size // mult * ngf * mult,
ngf * mult,
bias=False),
nn.ReLU(True),
nn.Linear(ngf * mult, ngf * mult, bias=False),
nn.ReLU(True)
]
self.gamma = nn.Linear(ngf * mult, ngf * mult, bias=False)
self.beta = nn.Linear(ngf * mult, ngf * mult, bias=False)
# Up-Sampling Bottleneck
for i in range(n_blocks):
setattr(self, 'UpBlock1_' + str(i + 1),
ResnetAdaILNBlock(ngf * mult, use_bias=False))
# Up-Sampling
UpBlock2 = []
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
UpBlock2 += [
nn.Upsample(scale_factor=2, mode='nearest'),
nn.ReflectionPad2d(1),
nn.Conv2d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=1,
padding=0,
bias=False),
ILN(int(ngf * mult / 2)),
nn.ReLU(True)
]
UpBlock2 += [
nn.ReflectionPad2d(3),
nn.Conv2d(ngf,
output_nc,
kernel_size=7,
stride=1,
padding=0,
bias=False),
nn.Tanh()
]
self.DownBlock = nn.Sequential(*DownBlock)
self.FC = nn.Sequential(*FC)
self.UpBlock2 = nn.Sequential(*UpBlock2)
