def __init__(self, fin, fout, opt, use_se=False, dilation=1):
super().__init__()
# Attributes
self.learned_shortcut = (fin != fout)
fmiddle = min(fin, fout)
self.opt = opt
self.pad_type = 'nozero'
self.use_se = use_se
# create conv layers
if self.pad_type != 'zero':
self.pad = nn.ReflectionPad2d(dilation)
self.conv_0 = nn.Conv2d(fin,
fmiddle,
kernel_size=3,
padding=0,
dilation=dilation)
self.conv_1 = nn.Conv2d(fmiddle,
fout,
kernel_size=3,
padding=0,
dilation=dilation)
else:
self.conv_0 = nn.Conv2d(fin,
fmiddle,
kernel_size=3,
padding=dilation,
dilation=dilation)
self.conv_1 = nn.Conv2d(fmiddle,
fout,
kernel_size=3,
padding=dilation,
dilation=dilation)
if self.learned_shortcut:
self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
# apply spectral norm if specified
if 'spectral' in opt.norm_G:
if opt.eqlr_sn:
self.conv_0 = equal_lr(self.conv_0)
self.conv_1 = equal_lr(self.conv_1)
if self.learned_shortcut:
self.conv_s = equal_lr(self.conv_s)
else:
self.conv_0 = spectral_norm(self.conv_0)
self.conv_1 = spectral_norm(self.conv_1)
if self.learned_shortcut:
self.conv_s = spectral_norm(self.conv_s)
# define normalization layers
spade_config_str = opt.norm_G.replace('spectral', '')
if 'spade_ic' in opt:
ic = opt.spade_ic
else:
ic = 0 + (3 if 'warp' in opt.CBN_intype else
0) + (opt.semantic_nc if 'mask' in opt.CBN_intype else 0)
self.norm_0 = SPADE(spade_config_str,
fin,
ic,
PONO=opt.PONO,
use_apex=opt.apex)
self.norm_1 = SPADE(spade_config_str,
fmiddle,
ic,
PONO=opt.PONO,
use_apex=opt.apex)
if self.learned_shortcut:
self.norm_s = SPADE(spade_config_str,
fin,
ic,
PONO=opt.PONO,
use_apex=opt.apex)
if use_se:
self.se_layar = SELayer(fout)
def snlinear(in_features, out_features):
return spectral_norm(
nn.Linear(in_features=in_features, out_features=out_features))
def __init__(self): super(DCGAN_D, self).__init__() self.dense = torch.nn.Linear(512 * 4 * 4, 1) if param.spectral: model = [spectral_norm(torch.nn.Conv2d(param.n_colors*2, 64, kernel_size=3, stride=1, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(64, 64, kernel_size=4, stride=2, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(256, 256, kernel_size=4, stride=2, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True), spectral_norm(torch.nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True)), torch.nn.LeakyReLU(0.1, inplace=True)] else: model = [torch.nn.Conv2d(param.n_colors*2, 64, kernel_size=3, stride=1, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(64)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(64, 64, kernel_size=4, stride=2, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(64)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(128)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(128)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(256)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(256, 256, kernel_size=4, stride=2, padding=1, bias=True)] if not param.no_batch_norm_D: model += [torch.nn.BatchNorm2d(256)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] model += [torch.nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True)] if param.Tanh_GD: model += [torch.nn.Tanh()] else: model += [torch.nn.LeakyReLU(0.1, inplace=True)] self.model = torch.nn.Sequential(*model) self.sig = torch.nn.Sigmoid()
def __init__(self, feature_dim=2048, adaptor_dim=256):
super(DomainAdaptor, self).__init__()
## initialization
self.bottleneck_layer = spectral_norm(nn.Linear(feature_dim, adaptor_dim))
self.bottleneck_layer.weight.data.normal_(0, 0.005)
self.bottleneck_layer.bias.data.fill_(0.1)
def __init__(self, inplanes, planes, kernel_size, padding, dilation):
super(_ASPPModule, self).__init__()
self.atrous_conv = spectral_norm(nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
stride=1, padding=padding, dilation=dilation, bias=False))
self.relu = nn.LeakyReLU(0.2)
def LinearSN(in_features, out_features, bias=True):
A = spectral_norm(nn.Linear(in_features, out_features, bias=bias))
# A.__class__.__name__ = 'LinearSN' ## [TODO]
return A
def upsampleLayer(inplanes, outplanes, upsample='basic', use_sn=True):
# padding_type = 'zero'
if upsample == 'basic' and not use_sn:
upconv = [
nn.ConvTranspose2d(inplanes,
outplanes,
kernel_size=3,
stride=2,
padding=1,
output_padding=1)
]
elif upsample == 'bilinear' and not use_sn:
upconv = [
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.ReflectionPad2d(1),
nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=1, padding=0)
]
elif upsample == 'nearest' and not use_sn:
upconv = [
nn.Upsample(scale_factor=2, mode='nearest'),
nn.ReflectionPad2d(1),
nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=1, padding=0)
]
elif upsample == 'subpixel' and not use_sn:
upconv = [
nn.Conv2d(inplanes,
outplanes * 4,
kernel_size=3,
stride=1,
padding=1),
nn.PixelShuffle(2)
]
elif upsample == 'basic' and use_sn:
upconv = [
spectral_norm(
nn.ConvTranspose2d(inplanes,
outplanes,
kernel_size=3,
stride=2,
padding=1,
output_padding=1))
]
elif upsample == 'bilinear' and use_sn:
upconv = [
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.ReflectionPad2d(1),
spectral_norm(
nn.Conv2d(inplanes,
outplanes,
kernel_size=3,
stride=1,
padding=0))
]
elif upsample == 'nearest' and use_sn:
upconv = [
nn.Upsample(scale_factor=2, mode='nearest'),
nn.ReflectionPad2d(1),
spectral_norm(
nn.Conv2d(inplanes,
outplanes,
kernel_size=3,
stride=1,
padding=0))
]
elif upsample == 'subpixel' and use_sn:
upconv = [
spectral_norm(
nn.Conv2d(inplanes,
outplanes * 4,
kernel_size=3,
stride=1,
padding=1)),
nn.PixelShuffle(2)
]
else:
raise NotImplementedError('upsample layer [%s] not implemented' %
upsample)
return upconv
def __init__(self,
opt,
input_nc=6,
ndf=32,
n_layers=0,
norm_layer=nn.BatchNorm2d,
use_sigmoid=True,
gpu_ids=[],
use_sn=False):
super(LabelChannelGatedResnetConvResnetD, self).__init__()
self.opt = opt
opt.nsalient = max(10, opt.n_classes)
self.label_embedding = nn.Embedding(opt.n_classes, opt.nsalient)
use_sn = opt.spectral_D
use_sigmoid = opt.no_lsgan
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
ndf = opt.ndf
kw = 4
padw = 1
sequence = []
nf_mult = 1
nf_mult_prev = 1
for n in range(1, n_layers):
nf_mult_prev = nf_mult
#nf_mult = min(2**n, 8)
if use_sn:
sequence += [
spectral_norm(
nn.Conv2d(ndf * nf_mult_prev,
ndf * nf_mult,
kernel_size=kw,
stride=2,
padding=padw,
bias=use_bias)),
get_norm(ndf * nf_mult, opt.norm_D, opt.num_groups),
nn.LeakyReLU(0.2, True)
]
else:
sequence += [
nn.Conv2d(ndf * nf_mult_prev,
ndf * nf_mult,
kernel_size=kw,
stride=2,
padding=padw,
bias=use_bias),
get_norm(ndf * nf_mult, opt.norm_D, opt.num_groups),
nn.LeakyReLU(0.2, True)
]
if use_sn:
sequence += [
spectral_norm(
nn.Conv2d(ndf * nf_mult,
opt.ndisc_out_filters,
kernel_size=kw,
stride=1,
padding=padw))
]
else:
sequence += [
nn.Conv2d(ndf * nf_mult,
opt.ndisc_out_filters,
kernel_size=kw,
stride=1,
padding=padw)
]
if use_sigmoid:
sequence += [nn.Sigmoid()]
self.main_latter = nn.Sequential(*sequence)
main_block = []
#Input is z going to series of rsidual blocks
# First layer to map to ndf channel
if opt.spectral_D:
main_block += [
spectral_norm(
nn.Conv2d(opt.input_nc + opt.output_nc,
opt.ndf,
kernel_size=3,
stride=1,
padding=1))
]
else:
main_block += [
nn.Conv2d(opt.input_nc + opt.output_nc,
opt.ndf,
kernel_size=3,
stride=1,
padding=1)
]
# Sets of residual blocks start
for i in range(3):
main_block += [
GatedConvResBlock(opt.ndf,
opt.ndf,
dropout=opt.dropout,
use_sn=opt.spectral_D,
norm_layer=opt.norm_D,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(opt.ndres_down):
main_block += [
DownGatedConvResBlock(opt.ndf,
opt.ndf,
dropout=opt.dropout_D,
use_sn=opt.spectral_D,
norm_layer=opt.norm_D,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(opt.ndres - opt.ndres_down - 3):
main_block += [
GatedConvResBlock(opt.ndf,
opt.ndf,
dropout=opt.dropout_D,
use_sn=opt.spectral_D,
norm_layer=opt.norm_D,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
self.main = nn.Sequential(*main_block)
gate_block = []
gate_block += [Reshape(-1, 1, opt.nsalient)]
gate_block += [
nn.Conv1d(1, opt.ngf_gate, kernel_size=3, stride=1, padding=1)
]
gate_block += [nn.ReLU()]
for i in range(opt.ndres_gate):
gate_block += [ResBlock1D(opt.ndf_gate, opt.dropout_gate)]
# state_size (opt.batchSize,opt.ndf_gate,opt.nsalient)
gate_block += [Reshape(-1, opt.ndf_gate * opt.nsalient)]
self.gate = nn.Sequential(*gate_block)
gate_block_mult = []
gate_block_mult += [
nn.Linear(opt.ndf_gate * opt.nsalient, opt.ndres * opt.ndf)
]
gate_block_mult += [nn.Sigmoid()]
self.gate_mult = nn.Sequential(*gate_block_mult)
if opt.gate_affine:
gate_block_add = []
gate_block_add += [
nn.Linear(opt.ndf_gate * opt.nsalient, opt.ndres * opt.ndf)
]
gate_block_add += [nn.Tanh()]
self.gate_add = nn.Sequential(*gate_block_add)
def __init__(self,
output_scale,
noise_size=120,
num_classes=0,
out_channels=3,
base_channels=96,
input_scale=4,
with_shared_embedding=True,
shared_dim=128,
sn_eps=1e-6,
sn_style='ajbrock',
init_type='ortho',
split_noise=True,
act_cfg=dict(type='ReLU'),
upsample_cfg=dict(type='nearest', scale_factor=2),
with_spectral_norm=True,
auto_sync_bn=True,
blocks_cfg=dict(type='BigGANGenResBlock'),
arch_cfg=None,
out_norm_cfg=dict(type='BN'),
pretrained=None,
rgb2bgr=False):
super().__init__()
self.noise_size = noise_size
self.num_classes = num_classes
self.shared_dim = shared_dim
self.with_shared_embedding = with_shared_embedding
self.output_scale = output_scale
self.arch = arch_cfg if arch_cfg else self._get_default_arch_cfg(
self.output_scale, base_channels)
self.input_scale = input_scale
self.split_noise = split_noise
self.blocks_cfg = deepcopy(blocks_cfg)
self.upsample_cfg = deepcopy(upsample_cfg)
self.rgb2bgr = rgb2bgr
self.sn_style = sn_style
# Validity Check
# If 'num_classes' equals to zero, we shall set 'with_shared_embedding'
# to False.
if num_classes == 0:
assert not self.with_shared_embedding
else:
if not self.with_shared_embedding:
# If not `with_shared_embedding`, we will use `nn.Embedding` to
# replace the original `Linear` layer in conditional BN.
# Meanwhile, we do not adopt split noises.
assert not self.split_noise
# If using split latents, we may need to adjust noise_size
if self.split_noise:
# Number of places z slots into
self.num_slots = len(self.arch['in_channels']) + 1
self.noise_chunk_size = self.noise_size // self.num_slots
# Recalculate latent dimensionality for even splitting into chunks
self.noise_size = self.noise_chunk_size * self.num_slots
else:
self.num_slots = 1
self.noise_chunk_size = 0
# First linear layer
self.noise2feat = nn.Linear(
self.noise_size // self.num_slots,
self.arch['in_channels'][0] * (self.input_scale**2))
if with_spectral_norm:
if sn_style == 'torch':
self.noise2feat = spectral_norm(self.noise2feat, eps=sn_eps)
elif sn_style == 'ajbrock':
self.noise2feat = SNLinear(self.noise_size // self.num_slots,
self.arch['in_channels'][0] *
(self.input_scale**2),
eps=sn_eps)
else:
raise NotImplementedError(f'Your {sn_style} is not supported')
# If using 'shared_embedding', we will get an unified embedding of
# label for all blocks. If not, we just pass the label to each
# block.
if with_shared_embedding:
self.shared_embedding = nn.Embedding(num_classes, shared_dim)
else:
self.shared_embedding = nn.Identity()
if num_classes > 0:
self.dim_after_concat = (self.shared_dim + self.noise_chunk_size
if self.with_shared_embedding else
self.num_classes)
else:
self.dim_after_concat = self.noise_chunk_size
self.blocks_cfg.update(
dict(dim_after_concat=self.dim_after_concat,
act_cfg=act_cfg,
sn_eps=sn_eps,
sn_style=sn_style,
input_is_label=(num_classes > 0)
and (not with_shared_embedding),
with_spectral_norm=with_spectral_norm,
auto_sync_bn=auto_sync_bn))
self.conv_blocks = nn.ModuleList()
for index, out_ch in enumerate(self.arch['out_channels']):
# change args to adapt to current block
self.blocks_cfg.update(
dict(in_channels=self.arch['in_channels'][index],
out_channels=out_ch,
upsample_cfg=self.upsample_cfg
if self.arch['upsample'][index] else None))
self.conv_blocks.append(build_module(self.blocks_cfg))
if self.arch['attention'][index]:
self.conv_blocks.append(
SelfAttentionBlock(out_ch,
with_spectral_norm=with_spectral_norm,
sn_eps=sn_eps,
sn_style=sn_style))
self.output_layer = SNConvModule(self.arch['out_channels'][-1],
out_channels,
kernel_size=3,
padding=1,
with_spectral_norm=with_spectral_norm,
spectral_norm_cfg=dict(
eps=sn_eps, sn_style=sn_style),
act_cfg=act_cfg,
norm_cfg=out_norm_cfg,
bias=True,
order=('norm', 'act', 'conv'))
self.init_weights(pretrained=pretrained, init_type=init_type)
def __init__(self, cfg):
super(DiscriminatorAdditiveGANRes, self).__init__()
self.activation = ACTIVATIONS[cfg.activation]
self.resdown1 = ResDownBlock(3,
64,
downsample=True,
first_block=True,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
# state size. (64) x 64 x 64
self.resdown2 = ResDownBlock(64,
128,
downsample=True,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
# state size. (128) x 32 x 32
self.resdown3 = ResDownBlock(128,
256,
downsample=True,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
extra_channels = 0
if cfg.use_fd:
if cfg.disc_img_conditioning == 'concat' or cfg.disc_img_conditioning == 'res':
extra_channels += 256
if cfg.conditioning == 'concat':
extra_channels += cfg.disc_cond_channels
# state size. (256) x 16 x 16
self.resdown4 = ResDownBlock(256 + extra_channels,
512,
downsample=True,
self_attn=cfg.self_attention,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
# state size. (512) x 8 x 8
self.resdown5 = ResDownBlock(512,
1024,
downsample=True,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
# state size. (1024) x 4 x 4
self.resdown6 = ResDownBlock(1024,
1024,
downsample=False,
activation=cfg.activation,
use_spectral_norm=cfg.disc_sn)
self.linear = torch.nn.Linear(1024, 1)
if cfg.disc_sn:
self.linear = spectral_norm(self.linear)
condition_dim = cfg.hidden_dim
if cfg.conditioning == 'projection':
self.condition_projector = nn.Sequential(
torch.nn.Linear(condition_dim, 1024),
nn.ReLU(),
torch.nn.Linear(1024, 1024),
)
elif cfg.conditioning == 'concat':
self.condition_projector = nn.Sequential(
torch.nn.Linear(condition_dim, 1024),
nn.ReLU(),
torch.nn.Linear(1024, cfg.disc_cond_channels),
)
if cfg.aux_reg > 0:
self.aux_objective = nn.Sequential(
torch.nn.Linear(1024, 256),
nn.ReLU(),
torch.nn.Linear(256, cfg.num_objects),
)
self.cfg = cfg
def __init__(self, opt):
super(StochasticLabelBetaChannelGatedResnetConvResnetG,
self).__init__()
self.opt = opt
opt.nsalient = max(10, opt.n_classes)
self.label_embedding = nn.Embedding(opt.n_classes, opt.nsalient)
self.main_initial = nn.Sequential(
nn.Conv2d(3, opt.ngf, kernel_size=3, stride=1, padding=1),
get_norm(opt.ngf, opt.norm_G, opt.num_groups), nn.ReLU(True))
self.label_noise = nn.Linear(opt.nz, opt.nsalient)
main_block = []
#Input is z going to series of rsidual blocks
# Sets of residual blocks start
for i in range(3):
main_block += [
GatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(opt.ngres_up_down):
main_block += [
DownGatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(int(opt.ngres / 2 - opt.ngres_up_down - 3)):
main_block += [
GatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(int(opt.ngres / 2 - opt.ngres_up_down - 3)):
main_block += [
GatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(opt.ngres_up_down):
main_block += [
UpGatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
for i in range(3):
main_block += [
GatedConvResBlock(opt.ngf,
opt.ngf,
dropout=opt.dropout_G,
use_sn=opt.spectral_G,
norm_layer=opt.norm_G,
num_groups=opt.num_groups,
res_op=opt.res_op)
]
# Final layer to map to 3 channel
if opt.spectral_G:
main_block += [
spectral_norm(
nn.Conv2d(opt.ngf,
opt.nc,
kernel_size=3,
stride=1,
padding=1))
]
else:
main_block += [
nn.Conv2d(opt.ngf, opt.nc, kernel_size=3, stride=1, padding=1)
]
main_block += [nn.Tanh()]
self.main = nn.Sequential(*main_block)
gate_block = []
gate_block += [Reshape(-1, 1, opt.nsalient)]
gate_block += [
nn.Conv1d(1, opt.ngf_gate, kernel_size=3, stride=1, padding=1)
]
gate_block += [nn.ReLU()]
for i in range(opt.ngres_gate):
gate_block += [ResBlock1D(opt.ngf_gate, opt.dropout_gate)]
# state size (opt.batchSize, opt.ngf_gate, opt.nsalient)
gate_block += [Reshape(-1, opt.ngf_gate * opt.nsalient)]
self.gate = nn.Sequential(*gate_block)
gate_block_mult = []
gate_block_mult += [
nn.Linear(opt.ngf_gate * opt.nsalient, opt.ngres * opt.ngf)
]
gate_block_mult += [nn.Sigmoid()]
self.gate_mult = nn.Sequential(*gate_block_mult)
gate_block_add = gate_block
gate_block_add += [
nn.Linear(opt.ngf_gate * opt.nsalient, opt.ngres * opt.ngf)
]
gate_block_add += [nn.Hardtanh()]
self.gate_add = nn.Sequential(*gate_block_add)
def __init__(
self,
latent_dim=100,
n_feature_maps=64,
n_class=2,
embedding_dim=8,
n_attributes=7,
):
super(Generator, self).__init__()
self.latent_dim = latent_dim
self.n_feature_maps = n_feature_maps
self.n_class = n_class
self.embedding_dim = embedding_dim
self.n_attributes = n_attributes
self.embeds = []
for i in range(self.n_attributes):
self.embeds.append(nn.Embedding(self.n_class, self.embedding_dim))
self.embeds = nn.ModuleList(self.embeds)
self.conv_blocks = nn.Sequential(
U.spectral_norm(
nn.ConvTranspose2d(
(self.latent_dim + self.n_attributes * self.embedding_dim),
self.n_feature_maps * 8,
4,
1,
0,
bias=False,
)
),
nn.BatchNorm2d(self.n_feature_maps * 8),
nn.ReLU(True),
U.spectral_norm(
nn.ConvTranspose2d(
self.n_feature_maps * 8,
self.n_feature_maps * 4,
4,
2,
1,
bias=False,
)
),
nn.BatchNorm2d(self.n_feature_maps * 4),
nn.ReLU(True),
U.spectral_norm(
nn.ConvTranspose2d(
self.n_feature_maps * 4,
self.n_feature_maps * 2,
4,
2,
1,
bias=False,
)
),
nn.BatchNorm2d(self.n_feature_maps * 2),
nn.ReLU(True),
U.spectral_norm(
nn.ConvTranspose2d(
self.n_feature_maps * 2, self.n_feature_maps, 4, 2, 1, bias=False
)
),
nn.BatchNorm2d(self.n_feature_maps),
nn.ReLU(True),
U.spectral_norm(
nn.ConvTranspose2d(self.n_feature_maps, 3, 4, 2, 1, bias=False)
),
nn.Tanh(),
)
def __init__(self, args):
super(SNGANDiscriminator, self).__init__()
self.args = args
self.nc = self.args.img_channels
self.ndf = self.args.disc_lat
cuda = True if torch.cuda.is_available() else False
self.this_device = torch.device("cuda" if cuda else "cpu")
if self.args.img_size == 64:
self.main = nn.Sequential(
# input is (nc) x 64 x 64
spectral_norm(nn.Conv2d(self.nc, self.ndf, 4, 2, 1,
bias=False)),
nn.BatchNorm2d(self.ndf),
nn.LeakyReLU(0.2, inplace=False),
# state size. (ndf) x 32 x 32
spectral_norm(
nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1, bias=False)),
nn.BatchNorm2d(self.ndf * 2),
nn.LeakyReLU(0.2, inplace=False),
# state size. (ndf*2) x 16 x 16
spectral_norm(
nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1,
bias=False)),
nn.BatchNorm2d(self.ndf * 4),
nn.LeakyReLU(0.2, inplace=False),
# state size. (ndf*4) x 8 x 8
spectral_norm(
nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1,
bias=False)),
nn.BatchNorm2d(self.ndf * 8),
nn.LeakyReLU(0.2, inplace=False),
# state size. (ndf*8) x 4 x 4
spectral_norm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0, bias=False)),
nn.Sigmoid())
elif self.args.img_size == 32:
# 64 uses kernel = 4, stride = 2, padding = 1. To reduce (64,64) to 1 via 5 layers.
# 32 uses kernel = 4, strude = 2, padding = 1. to reduce (32,32) to 1 via 4 layers.
# state size. (ndf) x 32 x 32
self.l21 = spectral_norm(
nn.Conv2d(self.nc, self.ndf * 2, 4, 2, 1, bias=False))
self.l22 = nn.BatchNorm2d(self.ndf * 2)
self.l23 = nn.LeakyReLU(0.2, inplace=False)
# state size. (ndf*2) x 16 x 16
self.l31 = spectral_norm(
nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1, bias=False))
self.l32 = nn.BatchNorm2d(self.ndf * 4)
self.l33 = nn.LeakyReLU(0.2, inplace=False)
# state size. (ndf*4) x 8 x 8
self.l41 = spectral_norm(
nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1, bias=False))
self.l42 = nn.BatchNorm2d(self.ndf * 8)
self.l43 = nn.LeakyReLU(0.2, inplace=False)
# state size. (ndf*8) x 4 x 4
self.l51 = spectral_norm(
nn.Conv2d(self.ndf * 8, 1, 4, 1, 0, bias=False))
self.l52 = nn.Sigmoid()
def __init__(self,
input_nc,
ndf=64,
n_layers=3,
norm_layer=nn.BatchNorm2d,
use_sigmoid=False,
getIntermFeat=False,
use_spectral_norm=False):
super(NLayerDiscriminator, self).__init__()
self.getIntermFeat = getIntermFeat
self.n_layers = n_layers
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
if use_spectral_norm:
sequence = [[
spectral_norm(
nn.Conv2d(input_nc,
ndf,
kernel_size=kw,
stride=2,
padding=padw)),
nn.LeakyReLU(0.2, True)
]]
else:
sequence = [[
nn.Conv2d(input_nc,
ndf,
kernel_size=kw,
stride=2,
padding=padw),
nn.LeakyReLU(0.2, True)
]]
nf = ndf
for n in range(1, n_layers):
nf_prev = nf
nf = min(nf * 2, 512)
if use_spectral_norm:
sequence += [[
spectral_norm(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=2,
padding=padw)),
norm_layer(nf),
nn.LeakyReLU(0.2, True)
]]
else:
sequence += [[
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=2,
padding=padw),
norm_layer(nf),
nn.LeakyReLU(0.2, True)
]]
nf_prev = nf
nf = min(nf * 2, 512)
if use_spectral_norm:
sequence += [[
spectral_norm(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=1,
padding=padw)),
norm_layer(nf),
nn.LeakyReLU(0.2, True)
]]
else:
sequence += [[
nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
norm_layer(nf),
nn.LeakyReLU(0.2, True)
]]
if use_spectral_norm:
sequence += [[
spectral_norm(
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw))
]]
else:
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
if use_sigmoid:
sequence += [[nn.Sigmoid()]]
if getIntermFeat:
for n in range(len(sequence)):
setattr(self, 'model' + str(n), nn.Sequential(*sequence[n]))
else:
sequence_stream = []
for n in range(len(sequence)):
sequence_stream += sequence[n]
self.model = nn.Sequential(*sequence_stream)
def __init__(self, base_acti_maps):
super(Discriminator, self).__init__()
self.lrelu_slope = 0.2
self.conv = nn.Sequential(
# input size: [B x 3 x 128 x 128]
spectral_norm(
nn.Conv2d(3, base_acti_maps * 2**0, 4, 2, 1, bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 0 x 64 x 64]
spectral_norm(
nn.Conv2d(base_acti_maps * 2**0,
base_acti_maps * 2**1,
4,
2,
1,
bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 1 x 32 x 32]
spectral_norm(
nn.Conv2d(base_acti_maps * 2**1,
base_acti_maps * 2**2,
4,
2,
1,
bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 2 x 16 x 16]
spectral_norm(
nn.Conv2d(base_acti_maps * 2**2,
base_acti_maps * 2**3,
4,
2,
1,
bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 3 x 8 x 8]
spectral_norm(
nn.Conv2d(base_acti_maps * 2**3,
base_acti_maps * 2**4,
4,
2,
1,
bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 4 x 4 x 4]
spectral_norm(
nn.Conv2d(base_acti_maps * 2**4,
base_acti_maps * 2**4,
4,
1,
0,
bias=False)),
nn.LeakyReLU(self.lrelu_slope, inplace=True),
# state size: [B x base * 2 ** 4 x 1 x 1]
# nn.AdaptiveAvgPool2d(1)
)
self.img_head = nn.Sequential(
spectral_norm(nn.Linear(base_acti_maps * 2**4 + 15, 1)))
self.hair_head = nn.Sequential(
spectral_norm(nn.Linear(base_acti_maps * 2**4, 6)), nn.Softmax(1))
self.eye_head = nn.Sequential(
spectral_norm(nn.Linear(base_acti_maps * 2**4, 4)), nn.Softmax(1))
self.face_head = nn.Sequential(
spectral_norm(nn.Linear(base_acti_maps * 2**4, 3)), nn.Softmax(1))
self.glass_head = nn.Sequential(
spectral_norm(nn.Linear(base_acti_maps * 2**4, 2)), nn.Softmax(1))
def __init__(self,
input_scale,
num_classes=0,
in_channels=3,
out_channels=1,
base_channels=96,
sn_eps=1e-6,
sn_style='ajbrock',
init_type='ortho',
act_cfg=dict(type='ReLU'),
with_spectral_norm=True,
blocks_cfg=dict(type='BigGANDiscResBlock'),
arch_cfg=None,
pretrained=None):
super().__init__()
self.num_classes = num_classes
self.out_channels = out_channels
self.input_scale = input_scale
self.in_channels = in_channels
self.base_channels = base_channels
self.arch = arch_cfg if arch_cfg else self._get_default_arch_cfg(
self.input_scale, self.in_channels, self.base_channels)
self.blocks_cfg = deepcopy(blocks_cfg)
self.blocks_cfg.update(
dict(act_cfg=act_cfg,
sn_eps=sn_eps,
sn_style=sn_style,
with_spectral_norm=with_spectral_norm))
self.sn_style = sn_style
self.conv_blocks = nn.ModuleList()
for index, out_ch in enumerate(self.arch['out_channels']):
# change args to adapt to current block
self.blocks_cfg.update(
dict(in_channels=self.arch['in_channels'][index],
out_channels=out_ch,
with_downsample=self.arch['downsample'][index],
is_head_block=(index == 0)))
self.conv_blocks.append(build_module(self.blocks_cfg))
if self.arch['attention'][index]:
self.conv_blocks.append(
SelfAttentionBlock(out_ch,
with_spectral_norm=with_spectral_norm,
sn_eps=sn_eps,
sn_style=sn_style))
self.activate = build_activation_layer(act_cfg)
self.decision = nn.Linear(self.arch['out_channels'][-1], out_channels)
if with_spectral_norm:
if sn_style == 'torch':
self.decision = spectral_norm(self.decision, eps=sn_eps)
elif sn_style == 'ajbrock':
self.decision = SNLinear(self.arch['out_channels'][-1],
out_channels,
eps=sn_eps)
else:
raise NotImplementedError('sn style')
if self.num_classes > 0:
self.proj_y = nn.Embedding(self.num_classes,
self.arch['out_channels'][-1])
if with_spectral_norm:
if sn_style == 'torch':
self.proj_y = spectral_norm(self.proj_y, eps=sn_eps)
elif sn_style == 'ajbrock':
self.proj_y = SNEmbedding(self.num_classes,
self.arch['out_channels'][-1],
eps=sn_eps)
else:
raise NotImplementedError('sn style')
self.init_weights(pretrained=pretrained, init_type=init_type)
def snlinear(in_features, out_features, bias=True):
return spectral_norm(
nn.Linear(in_features=in_features,
out_features=out_features,
bias=bias))
def __init__(self,
dim_z=128,
img_c=3,
fm_base=64,
bottom_width=4,
n_classes=0,
bn=False,
sn=False,
sa=False):
super(SN_ResNet_ImgNet128_Gen, self).__init__()
self.bn = bn
self.bottom_width = bottom_width
self.activation = F.relu
self.dim_z = dim_z
self.n_classes = n_classes
self.sa = sa
self.l1 = nn.Linear(dim_z, (bottom_width**2) * fm_base * 16)
nn.init.xavier_uniform_(self.l1.weight)
self.block2 = UpBlock(fm_base * 16,
fm_base * 16,
activation=self.activation,
upsample=True,
n_classes=n_classes,
bn=bn,
sn=sn)
self.block3 = UpBlock(fm_base * 16,
fm_base * 8,
activation=self.activation,
upsample=True,
n_classes=n_classes,
bn=bn,
sn=sn)
self.block4 = UpBlock(fm_base * 8,
fm_base * 4,
activation=self.activation,
upsample=True,
n_classes=n_classes,
bn=bn,
sn=sn)
if self.sa:
self.attn1 = Attention(ch=fm_base * 4, sn=sn)
self.block5 = UpBlock(fm_base * 4,
fm_base * 2,
activation=self.activation,
upsample=True,
n_classes=n_classes,
bn=bn,
sn=sn)
self.block6 = UpBlock(fm_base * 2,
fm_base * 1,
activation=self.activation,
upsample=True,
n_classes=n_classes,
bn=bn,
sn=sn)
if bn:
self.finalbn = nn.BatchNorm2d(fm_base)
self.l6 = nn.Conv2d(fm_base, img_c, kernel_size=3, stride=1, padding=1)
nn.init.xavier_uniform_(self.l6.weight)
if sn:
from torch.nn.utils import spectral_norm
self.l1 = spectral_norm(self.l1)
self.l6 = spectral_norm(self.l6)
def __init__(self, args):
super(Discriminator, self).__init__()
self.name = f'discriminator_{args.dataset}'
self.bias = False
channels = 32
layers = [
nn.Conv2d(3,
channels,
kernel_size=3,
stride=1,
padding=1,
bias=self.bias),
nn.LeakyReLU(0.2, True)
]
for i in range(args.d_layers):
layers += [
nn.Conv2d(channels,
channels * 2,
kernel_size=3,
stride=2,
padding=1,
bias=self.bias),
nn.LeakyReLU(0.2, True),
nn.Conv2d(channels * 2,
channels * 4,
kernel_size=3,
stride=1,
padding=1,
bias=self.bias),
nn.InstanceNorm2d(channels * 4),
nn.LeakyReLU(0.2, True),
]
channels *= 4
layers += [
nn.Conv2d(channels,
channels,
kernel_size=3,
stride=1,
padding=1,
bias=self.bias),
nn.InstanceNorm2d(channels),
nn.LeakyReLU(0.2, True),
nn.Conv2d(channels,
1,
kernel_size=3,
stride=1,
padding=1,
bias=self.bias),
]
if args.use_sn:
for i in range(len(layers)):
if isinstance(layers[i], nn.Conv2d):
layers[i] = spectral_norm(layers[i])
self.discriminate = nn.Sequential(*layers)
initialize_weights(self)
def __init__(self,
img_c=3,
fm_base=64,
n_classes=0,
bn=False,
sn=False,
sa=False,
c_metric='',
wide=False):
super(SN_ResNet_ImgNet128_Dis, self).__init__()
self.wide = wide
self.activation = F.relu
self.n_classes = n_classes
self.c_metric = c_metric
self.sa = sa
self.block1 = OptimizedBlock(img_c, fm_base * 1, bn=bn, sn=sn)
self.block2 = DownBlock(fm_base,
fm_base * 2,
(fm_base * 2 if self.wide else fm_base),
activation=self.activation,
downsample=True,
bn=bn,
sn=sn)
if self.sa:
self.attn1 = Attention(ch=fm_base * 2, sn=sn)
self.block3 = DownBlock(fm_base * 2,
fm_base * 4,
(fm_base * 4 if self.wide else fm_base * 2),
activation=self.activation,
downsample=True,
bn=bn,
sn=sn)
self.block4 = DownBlock(fm_base * 4,
fm_base * 8,
(fm_base * 8 if self.wide else fm_base * 4),
activation=self.activation,
downsample=True,
bn=bn,
sn=sn)
self.block5 = DownBlock(fm_base * 8,
fm_base * 16,
(fm_base * 16 if self.wide else fm_base * 8),
activation=self.activation,
downsample=True,
bn=bn,
sn=sn)
self.block6 = DownBlock(fm_base * 16,
fm_base * 16,
(fm_base * 16 if self.wide else fm_base * 16),
activation=self.activation,
downsample=False,
bn=bn,
sn=sn)
# bin
self.o_b = nn.Linear(fm_base * 16, 1)
nn.init.xavier_uniform_(self.o_b.weight)
# class
if n_classes > 0:
if c_metric in [Metric_Types[0], Metric_Types[2]]:
self.o_c = nn.Linear(fm_base * 16, n_classes)
elif c_metric == Metric_Types[1]:
self.o_c = nn.Embedding(n_classes, fm_base * 16)
nn.init.xavier_uniform_(self.o_c.weight)
if sn:
from torch.nn.utils import spectral_norm
self.o_b = spectral_norm(self.o_b)
if n_classes > 0 and (c_metric in Metric_Types):
self.o_c = spectral_norm(self.o_c)
def __init__(self, channel, ):
super().__init__()
self.conv = nn.Sequential( spectral_norm(nn.Conv2d(channel, channel, kernel_size=3, stride=1, padding=1, bias=False)),
nn.BatchNorm2d(channel),
nn.ReLU())
def add_sn(m):
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
return spectral_norm(m)
else:
return m
def _conv1d_spect(ni:int, no:int, ks:int=1, stride:int=1, padding:int=0, bias:bool=False):
"Create and initialize a `nn.Conv1d` layer with spectral normalization."
conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias)
nn.init.kaiming_normal_(conv.weight)
if bias: conv.bias.data.zero_()
return spectral_norm(conv)
def __init__(self,
n_hidden,
hidden_size,
activation_fn,
activation_slope,
init_method,
spect_norm=True,
norm='layer',
res_block=False,
dropout=False,
dropout_p=0.5):
super().__init__()
# Define activation function.
if activation_fn == 'relu':
activation = nn.ReLU(inplace=True)
elif activation_fn == 'leakyrelu':
activation = nn.LeakyReLU(inplace=True,
negative_slope=activation_slope)
elif activation_fn == 'tanh':
activation = nn.Tanh()
else:
raise NotImplementedError('Check activation_fn.')
if norm == 'layer':
norm = nn.LayerNorm
else:
norm = None
modules = [
spectral_norm(nn.Linear(1, hidden_size))
if spect_norm else nn.Linear(1, hidden_size)
]
if norm:
modules += [norm(hidden_size)]
for _ in range(n_hidden):
# Add dropout.
if dropout:
modules += [nn.Dropout(dropout_p)]
# Add act and layer.
if res_block:
modules += [
activation,
ResidualBlock(hidden_size, hidden_size, activation,
spect_norm, norm)
]
else:
modules += [
activation,
spectral_norm(nn.Linear(hidden_size, hidden_size))
if spect_norm else nn.Linear(hidden_size, hidden_size)
]
if norm:
modules += [norm(hidden_size)]
if dropout:
modules += [nn.Dropout(dropout_p)]
modules += [activation]
modules += [
spectral_norm(nn.Linear(hidden_size, 1))
if spect_norm else nn.Linear(hidden_size, 1)
]
self.model = nn.Sequential(*modules)
self.init_method = init_method
self.model.apply(self.__init)
def __init__(self, input_dim, output_dim, bias=True):
super(LinearSpec, self).__init__()
self.linear = spectral_norm(nn.Linear(input_dim, output_dim,
bias=bias))
def __init__(self,
submodule,
nfm=64,
in_frames=None,
out_frames=None,
outermost=False,
innermost=False):
super(U_Net_Block, self).__init__()
self.outermost = outermost
self.innermost = innermost
self.nfm = nfm
if outermost:
in_down_nc = in_frames * 3
out_up_nc = out_frames * 3
else:
in_down_nc = nfm
out_up_nc = nfm
downconv_1 = spectral_norm(nn.Conv2d(in_down_nc, in_down_nc, 3, 1, 1))
downnorm_1 = nn.BatchNorm2d(in_down_nc)
downconv = spectral_norm(nn.Conv2d(in_down_nc, nfm * 2, 3, 2, 1))
downnorm = nn.BatchNorm2d(nfm * 2)
upconv_1 = nn.ConvTranspose2d(nfm * 4, nfm * 4, 3, 1, 1)
upnorm_1 = nn.BatchNorm2d(nfm * 4)
upconv = nn.ConvTranspose2d(nfm * 4, out_up_nc * 2, 3, 2, 0)
upnorm = nn.BatchNorm2d(out_up_nc * 2)
relu = nn.ReLU()
tanh = nn.Tanh()
if outermost:
upconv_1 = nn.ConvTranspose2d(nfm * 4, nfm * 4, 3, 1, 1)
upnorm = nn.BatchNorm2d(nfm * 4)
upconv = nn.ConvTranspose2d(nfm * 4, out_up_nc, 3, 2, 1)
down = [downconv_1, downnorm_1, relu, downconv, downnorm, relu]
up = [upconv_1, upnorm, relu, upconv, tanh]
elif innermost:
res_nfm = 64
downconv_1 = spectral_norm(nn.Conv2d(nfm, nfm, 3, 1, 1))
downnorm_1 = nn.BatchNorm2d(nfm)
downconv = spectral_norm(nn.Conv2d(nfm, res_nfm, 3, 2, 1))
downnorm = nn.BatchNorm2d(res_nfm)
upconv_1 = nn.ConvTranspose2d(res_nfm, res_nfm, 3, 1, 1)
upnorm_1 = nn.BatchNorm2d(res_nfm)
upconv = nn.ConvTranspose2d(res_nfm, out_up_nc, 3, 2, 0)
upnorm = nn.BatchNorm2d(out_up_nc)
down = [downconv_1, downnorm_1, relu, downconv, downnorm, relu]
up = [upconv_1, upnorm_1, relu, upconv, upnorm, relu]
else:
downconv_1 = spectral_norm(
nn.Conv2d(in_down_nc * 2, in_down_nc * 2, 3, 1, 1))
downnorm_1 = nn.BatchNorm2d(in_down_nc * 2)
downconv = spectral_norm(
nn.Conv2d(in_down_nc * 2, nfm * 2, 3, 2, 1))
down = [downconv_1, downnorm_1, relu, downconv, downnorm, relu]
up = [upconv_1, upnorm_1, relu, upconv, upnorm, relu]
self.u_block = down + [submodule] + up
self.u_block = nn.Sequential(*self.u_block)
def sn_embedding(num_embeddings, embedding_dim):
return spectral_norm(
nn.Embedding(num_embeddings=num_embeddings,
embedding_dim=embedding_dim))
def __init__(self, reduced=False, spectral_norm=False, batch_norm=True):
super(Discriminator, self).__init__()
if reduced == False:
self.conv_channels = [(2, 32), (32, 64), (64, 64), (64, 128),
(128, 128), (128, 256), (256, 256),
(256, 512), (512, 512), (512, 1024),
(1024, 2048)]
self.conv_params = {
'kernel_size': 31,
'stride': 2,
'padding': 15,
'dilation': 1,
'groups': 1,
'bias': False
}
else:
self.conv_channels = [(2, 64), (64, 128), (128, 256), (256, 512),
(512, 1024)]
self.conv_params = {
'kernel_size': 31,
'stride': 4,
'padding': 15,
'dilation': 1,
'groups': 1,
'bias': False
}
self.Conv_layers = []
"""
Conv_layers:
[B x 2 x 16384]->[B x 16 x 8192]
[B x 16 x 8192]->[B x 32 x 4096]
[B x 32 x 4096]->[B x 32 x 2048]
[B x 32 x 2048]->[B x 64 x 1024]
[B x 64 x 1024]->[B x 64 x 512]
[B x 64 x 512]->[B x 128 x 256]
[B x 128 x 256]->[B x 128 x 128]
[B x 128 x 128]->[B x 256 x 64]
[B x 256 x 64]->[B x 256 x 32]
[B x 256 x 32]->[B x 512 x 16]
[B x 512 x 16]->[B x 1024 x 8]
"""
for layer_idx in range(len(self.conv_channels)):
self.conv_params['in_channels'] = self.conv_channels[layer_idx][0]
self.conv_params['out_channels'] = self.conv_channels[layer_idx][1]
if batch_norm == True:
modules = [
nn.Conv1d(**self.conv_params),
nn.BatchNorm1d(self.conv_params['out_channels']),
nn.PReLU(),
nn.Dropout()
]
else:
modules = [
nn.Conv1d(**self.conv_params),
nn.PReLU(),
nn.Dropout()
]
# if (layer_idx + 1) % 3 == 0:
# modules.insert(1, nn.Dropout())
self.Conv_layers.extend(modules)
self.Conv_layers = nn.Sequential(*self.Conv_layers)
self.Linear_layers = nn.Sequential(nn.Linear(16384, 256), nn.PReLU(),
nn.Linear(256, 128), nn.PReLU(),
nn.Linear(128, 1))
weights_init(self)
if spectral_norm == True:
for m in self.modules():
if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
m = utils.spectral_norm(m)
def __init__(self, fin, fout, opt, Block_Name=None, use_rgb=True):
super().__init__()
# sean switch
if opt.norm_mode != 'sean':
self.use_sean = False
elif opt.norm_mode == 'sean':
self.use_sean = True
# SEAN/ ACE-block specific attributes
self.use_rgb = use_rgb
self.Block_Name = Block_Name
self.status = opt.status
# Attributes
self.learned_shortcut = (fin != fout)
fmiddle = min(fin, fout)
# create conv layers
# There is some weird line here that is only relevant when instance maps are used.
#
# add_channels = 1 if (opt.norm_mode == 'clade' and not opt.no_instance) else 0
self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
if self.learned_shortcut:
self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
# apply spectral norm if specified
if 'spectral' in opt.norm_G:
self.conv_0 = spectral_norm(self.conv_0)
self.conv_1 = spectral_norm(self.conv_1)
if self.learned_shortcut:
self.conv_s = spectral_norm(self.conv_s)
# define normalization layers
#
# Mike: Added the if else option of clades architecture.py
#
# Added the option to choose sean blocks
spade_config_str = opt.norm_G.replace('spectral', '')
if opt.norm_mode == 'spade':
self.norm_0 = SPADE(spade_config_str, fin, opt.semantic_nc)
self.norm_1 = SPADE(spade_config_str, fmiddle, opt.semantic_nc)
if self.learned_shortcut:
self.norm_s = SPADE(spade_config_str, fin, opt.semantic_nc)
elif opt.norm_mode == 'clade':
input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0)
self.norm_0 = SPADELight(spade_config_str, fin, input_nc, opt.no_instance, opt.add_dist)
self.norm_1 = SPADELight(spade_config_str, fmiddle, input_nc, opt.no_instance, opt.add_dist)
if self.learned_shortcut:
self.norm_s = SPADELight(spade_config_str, fin, input_nc, opt.no_instance, opt.add_dist)
elif opt.norm_mode == 'sean':
self.use_sean = True
our_norm_type = 'spadesyncbatch3x3'
self.ace_0 = ACE(our_norm_type, fin, 3, ACE_Name= Block_Name + '_ACE_0', status=self.status, spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb)
self.ace_1 = ACE(our_norm_type, fmiddle, 3, ACE_Name= Block_Name + '_ACE_1', status=self.status, spade_params=[spade_config_str, fmiddle, opt.semantic_nc], use_rgb=use_rgb)
if self.learned_shortcut:
self.ace_s = ACE(our_norm_type, fin, 3, ACE_Name= Block_Name + '_ACE_s', status=self.status, spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb)
else:
raise ValueError('%s is not a defined normalization method' % opt.norm_mode)
def __init__(self, in_channel):
super(SelfAttention, self).__init__()
self.query = spectral_norm(nn.Conv1d(in_channel, in_channel // 8, 1))
self.key = spectral_norm(nn.Conv1d(in_channel, in_channel // 8, 1))
self.value = spectral_norm(nn.Conv1d(in_channel, in_channel, 1))
self.gamma = nn.Parameter(torch.tensor(0.0))
