def __init__(self,
dim_z,
skip_init=False,
E_lr=2e-4,
adam_eps=1e-8,
E_B1=0.0,
E_B2=0.999,
E_mixed_precision=False,
name=None,
**kwargs):
super(Encoder_rd, self).__init__(Bottleneck, [3, 4, 6, 3],
width_per_group=64 * 16)
self.resblock1 = ResBlock(512 * 4, dim_z * 2)
self.resblock2 = ResBlock(dim_z * 2, dim_z * 2)
if not skip_init:
self.init_weights()
# Set name for saving and loading weights
self.name = name if name is not None else "Encoder"
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = E_lr, E_B1, E_B2, adam_eps
if E_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
def __init__(self, dim_z, I_depth, is_reverse=True, I_lr=2e-4,
adam_eps=1e-8, I_B1=0.0, I_B2=0.999, I_mixed_precision=False, name=None, **kwargs):
"""
Invertible network.
:param dim_z: Must be int.
:param depth: How many coupling layer are used.
:param is_reverse:
"""
super(Invert, self).__init__()
self._z_dim = dim_z
self._depth = I_depth
self._is_reverse = is_reverse
self.steps = []
self.steps = nn.ModuleList([step(self._z_dim, is_reverse) for _ in range(self._depth)])
# Set name, used for load and save weights
self.name = name if name is not None else 'Invert'
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = I_lr, I_B1, I_B2, adam_eps
if I_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
def __init__(self,
dim_z,
skip_init=False,
E_lr=2e-4,
E_init='ortho',
widden_factor=2,
adam_eps=1e-8,
E_B1=0.0,
E_B2=0.999,
E_mixed_precision=False,
name=None,
**kwargs):
super(Encoder_inv, self).__init__(Bottleneck, [3, 4, 6, 3],
width_per_group=64 * widden_factor)
self.downsample = nn.Sequential(
nn.Conv2d(512 * 4, dim_z, kernel_size=1, stride=1, bias=False),
nn.BatchNorm2d(dim_z))
self.out_layer = nn.Sequential(nn.Linear(dim_z, dim_z),
nn.ReLU(inplace=True))
self.init = E_init
if not skip_init:
self.init_weights()
# Set name used for save and load weights
self.name = name if name is not None else "Encoder"
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = E_lr, E_B1, E_B2, adam_eps
if E_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
def __init__(self,
dim_z,
skip_init=False,
E_lr=2e-4,
E_init='ortho',
widden_factor=2,
adam_eps=1e-8,
E_B1=0.0,
E_B2=0.999,
E_mixed_precision=False,
name=None,
**kwargs):
super(Encoder_out_layer, self).__init__()
self.dense = nn.Linear(2048, dim_z)
self.init = E_init
if not skip_init:
self.init_weights()
# Set name used for save and load weights
self.name = name if name is not None else "Encoder_out_layer"
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = E_lr, E_B1, E_B2, adam_eps
if E_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
def __init__(self,
dim_z=120,
L_depth=4,
skip_init=False,
L_init='ortho',
L_lr=2e-4,
adam_eps=1e-8,
L_B1=0.0,
L_B2=0.999,
L_mixed_precision=False,
name=None,
**kwargs):
super(LatentBinder, self).__init__()
self.layers = torch.nn.ModuleList([ResBlock() for _ in range(L_depth)])
self.out = torch.nn.Sequential(torch.nn.Linear(dim_z, 1),
torch.nn.ReLU(inplace=True))
self.init = L_init
if not skip_init:
self.init_weights()
self.name = name if name is not None else "LatentBinder"
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = L_lr, L_B1, L_B2, adam_eps
if L_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
def __init__(self,
G_ch=64,
dim_z=128,
bottom_width=4,
resolution=128,
G_kernel_size=3,
G_attn='64',
n_classes=1000,
num_G_SVs=1,
num_G_SV_itrs=1,
G_shared=True,
shared_dim=0,
hier=False,
cross_replica=False,
mybn=False,
G_activation=nn.ReLU(inplace=False),
G_lr=5e-5,
G_B1=0.0,
G_B2=0.999,
adam_eps=1e-8,
BN_eps=1e-5,
SN_eps=1e-12,
G_mixed_precision=False,
G_fp16=False,
G_init='ortho',
skip_init=False,
no_optim=False,
G_param='SN',
norm_style='bn',
**kwargs):
super(Generator, self).__init__()
# Channel width mulitplier
self.ch = G_ch
# Dimensionality of the latent space
self.dim_z = dim_z
# The initial spatial dimensions
self.bottom_width = bottom_width
# Resolution of the output
self.resolution = resolution
# Kernel size?
self.kernel_size = G_kernel_size
# Attention?
self.attention = G_attn
# number of classes, for use in categorical conditional generation
self.n_classes = n_classes
# Use shared embeddings?
self.G_shared = G_shared
# Dimensionality of the shared embedding? Unused if not using G_shared
self.shared_dim = shared_dim if shared_dim > 0 else dim_z
# Hierarchical latent space?
self.hier = hier
# Cross replica batchnorm?
self.cross_replica = cross_replica
# Use my batchnorm?
self.mybn = mybn
# nonlinearity for residual blocks
self.activation = G_activation
# Initialization style
self.init = G_init
# Parameterization style
self.G_param = G_param
# Normalization style
self.norm_style = norm_style
# Epsilon for BatchNorm?
self.BN_eps = BN_eps
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# fp16?
self.fp16 = G_fp16
# Architecture dict
self.arch = G_arch(self.ch, self.attention)[resolution]
# If using hierarchical latents, adjust z
if self.hier:
# Number of places z slots into
self.num_slots = len(self.arch['in_channels']) + 1
self.z_chunk_size = (self.dim_z // self.num_slots)
# Recalculate latent dimensionality for even splitting into chunks
self.dim_z = self.z_chunk_size * self.num_slots
else:
self.num_slots = 1
self.z_chunk_size = 0
# Which convs, batchnorms, and linear layers to use
if self.G_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3,
padding=1,
num_svs=num_G_SVs,
num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_G_SVs,
num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
else:
self.which_conv = functools.partial(nn.Conv2d,
kernel_size=3,
padding=1)
self.which_linear = nn.Linear
# We use a non-spectral-normed embedding here regardless;
# For some reason applying SN to G's embedding seems to randomly cripple G
self.which_embedding = nn.Embedding
bn_linear = (functools.partial(self.which_linear, bias=False)
if self.G_shared else self.which_embedding)
self.which_bn = functools.partial(
layers.ccbn,
which_linear=bn_linear,
cross_replica=self.cross_replica,
mybn=self.mybn,
input_size=(self.shared_dim + self.z_chunk_size
if self.G_shared else self.n_classes),
norm_style=self.norm_style,
eps=self.BN_eps)
# Prepare model
# If not using shared embeddings, self.shared is just a passthrough
self.shared = (self.which_embedding(n_classes, self.shared_dim)
if G_shared else layers.identity())
# First linear layer
self.linear = self.which_linear(
self.dim_z // self.num_slots,
self.arch['in_channels'][0] * (self.bottom_width**2))
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
# while the inner loop is over a given block
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[
layers.GBlock(
in_channels=self.arch['in_channels'][index],
out_channels=self.arch['out_channels'][index],
which_conv=self.which_conv,
which_bn=self.which_bn,
activation=self.activation,
upsample=(functools.partial(F.interpolate, scale_factor=2)
if self.arch['upsample'][index] else None))
]]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in G at resolution %d' %
self.arch['resolution'][index])
self.blocks[-1] += [
layers.Attention(self.arch['out_channels'][index],
self.which_conv)
]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList(
[nn.ModuleList(block) for block in self.blocks])
# output layer: batchnorm-relu-conv.
# Consider using a non-spectral conv here
self.output_layer = nn.Sequential(
layers.bn(self.arch['out_channels'][-1],
cross_replica=self.cross_replica,
mybn=self.mybn), self.activation,
self.which_conv(self.arch['out_channels'][-1], 3))
# Initialize weights. Optionally skip init for testing.
if not skip_init:
self.init_weights()
# Set up optimizer
# If this is an EMA copy, no need for an optim, so just return now
if no_optim:
return
self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps
if G_mixed_precision:
print('Using fp16 adam in G...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
def __init__(self,
D_ch=64,
D_wide=True,
resolution=128,
D_kernel_size=3,
D_attn='64',
n_classes=1000,
num_D_SVs=1,
num_D_SV_itrs=1,
D_activation=nn.ReLU(inplace=False),
D_lr=2e-4,
D_B1=0.0,
D_B2=0.999,
adam_eps=1e-8,
SN_eps=1e-12,
output_dim=1,
D_mixed_precision=False,
D_fp16=False,
D_init='ortho',
skip_init=False,
D_param='SN',
**kwargs):
super(Discriminator, self).__init__()
# Width multiplier
self.ch = D_ch
# Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?
self.D_wide = D_wide
# Resolution
self.resolution = resolution
# Kernel size
self.kernel_size = D_kernel_size
# Attention?
self.attention = D_attn
# Number of classes
self.n_classes = n_classes
# Activation
self.activation = D_activation
# Initialization style
self.init = D_init
# Parameterization style
self.D_param = D_param
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# Fp16?
self.fp16 = D_fp16
# Architecture
self.arch = D_arch(self.ch, self.attention)[resolution]
# Which convs, batchnorms, and linear layers to use
# No option to turn off SN in D right now
if self.D_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3,
padding=1,
num_svs=num_D_SVs,
num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_D_SVs,
num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_embedding = functools.partial(layers.SNEmbedding,
num_svs=num_D_SVs,
num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
# Prepare model
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[
layers.DBlock(
in_channels=self.arch['in_channels'][index],
out_channels=self.arch['out_channels'][index],
which_conv=self.which_conv,
wide=self.D_wide,
activation=self.activation,
preactivation=(index > 0),
downsample=(nn.AvgPool2d(2)
if self.arch['downsample'][index] else None))
]]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in D at resolution %d' %
self.arch['resolution'][index])
self.blocks[-1] += [
layers.Attention(self.arch['out_channels'][index],
self.which_conv)
]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList(
[nn.ModuleList(block) for block in self.blocks])
# Linear output layer. The output dimension is typically 1, but may be
# larger if we're e.g. turning this into a VAE with an inference output
self.linear = self.which_linear(self.arch['out_channels'][-1],
output_dim)
# Embedding for projection discrimination
self.embed = self.which_embedding(self.n_classes,
self.arch['out_channels'][-1])
# Initialize weights
if not skip_init:
self.init_weights()
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps
if D_mixed_precision:
print('Using fp16 adam in D...')
from Utils import utils
self.optim = utils.Adam16(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(),
lr=self.lr,
betas=(self.B1, self.B2),
weight_decay=0,
eps=self.adam_eps)
