def __init__(self, in_channel, channel, n_res_block, n_res_channel,
embed_dim, norm, use_spectral_norm, out_activation):
super().__init__()
self.enc_b = ResEncoder(in_channel,
channel,
n_res_channel,
norm,
use_spectral_norm,
n_res_block,
out_activation,
stride=4)
self.enc_t = ResEncoder(channel,
channel,
n_res_channel,
norm,
use_spectral_norm,
n_res_block,
out_activation,
stride=2)
self.quantize_conv_t = norm_layer(nn.Conv2d(channel, embed_dim, 1),
norm, use_spectral_norm)
self.dec_t = ResDecoder(embed_dim,
channel,
n_res_channel,
embed_dim,
norm,
use_spectral_norm,
n_res_block,
out_activation="",
stride=2)
self.quantize_conv_b = norm_layer(
nn.Conv2d(embed_dim + channel, embed_dim, 1), norm,
use_spectral_norm)
def __init__(self, dim, dim_res, norm="BN", use_spectral_norm=False):
super().__init__()
self.block = nn.Sequential(
nn.ReLU(True),
norm_layer(nn.Conv2d(dim, dim_res, 3, 1, 1),
norm,
use_spectral_norm=use_spectral_norm),
nn.ReLU(True),
norm_layer(nn.Conv2d(dim_res, dim, 1),
norm,
use_spectral_norm=use_spectral_norm),
)
def __init__(self, in_channels, nf, out_channels, norm, use_spectral_norm,
n_layers, out_activation):
super().__init__()
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
self.layers = [
norm_layer(nn.Conv2d(in_channels, nf, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm),
nn.LeakyReLU(0.2, True)
]
kp = nf
for i in range(n_layers):
k = nf << i
self.layers.append(
norm_layer(nn.Conv2d(kp, k, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
self.layers.append(nn.LeakyReLU(0.2, True))
self.layers.append(
norm_layer(nn.Conv2d(k, k, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
self.layers.append(nn.LeakyReLU(0.2, True))
self.layers.append(nn.AvgPool2d(2))
kp = k
k = nf << n_layers
self.layers.append(
norm_layer(nn.Conv2d(kp, k, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
self.layers.append(nn.LeakyReLU(0.2, True))
self.layers.append(
norm_layer(nn.Conv2d(k, out_channels, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
if out_activation == "":
pass
elif out_activation == "sigmoid":
self.layers.append(nn.Sigmoid())
elif out_activation == "tanh":
self.layers.append(nn.Tanh())
else:
raise ValueError
self.layers = nn.Sequential(*self.layers)
def __init__(self, in_channels, nf, res_channels, norm, use_spectral_norm,
n_layers, out_activation, stride):
super().__init__()
if stride == 4:
self.layers = [
norm_layer(nn.Conv2d(in_channels, nf // 2, 4, 2, 1),
norm,
use_spectral_norm=use_spectral_norm),
nn.ReLU(True),
norm_layer(nn.Conv2d(nf // 2, nf, 4, 2, 1),
norm,
use_spectral_norm=use_spectral_norm),
nn.ReLU(True),
norm_layer(nn.Conv2d(nf, nf, 3, 1, 1),
norm,
use_spectral_norm=use_spectral_norm),
]
elif stride == 2:
self.layers = [
norm_layer(nn.Conv2d(in_channels, nf // 2, 4, 2, 1),
norm,
use_spectral_norm=use_spectral_norm),
nn.ReLU(True),
norm_layer(nn.Conv2d(nf // 2, nf, 3, 1, 1),
norm,
use_spectral_norm=use_spectral_norm),
]
else:
raise ValueError
for i in range(n_layers):
self.layers.append(ResBlock(nf, res_channels, norm))
if out_activation == "":
pass
elif out_activation == "sigmoid":
self.layers.append(nn.Sigmoid())
elif out_activation == "relu":
self.layers.append(nn.ReLU(inplace=True))
elif out_activation == "tanh":
self.layers.append(nn.Tanh())
else:
raise ValueError
self.layers = nn.Sequential(*self.layers)
def __init__(self, in_channels, nf, out_channels, norm, use_spectral_norm,
n_layers, out_activation):
super().__init__()
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
self.layers = []
kp = in_channels
for scale in range(n_layers - 1, -1, -1):
k = nf << scale
self.layers.append(
norm_layer(nn.Conv2d(kp, k, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
self.layers.append(nn.LeakyReLU(0.2, True))
self.layers.append(
norm_layer(nn.Conv2d(k, k, kw, stride=1, padding=pw),
norm,
use_spectral_norm=use_spectral_norm))
self.layers.append(nn.LeakyReLU(0.2, True))
self.layers.append(nn.Upsample(scale_factor=2))
kp = k
self.layers.append(
norm_layer(nn.Conv2d(kp, nf, kw, stride=1, padding=pw),
norm="",
use_spectral_norm=use_spectral_norm))
self.layers.append(
norm_layer(nn.Conv2d(kp, out_channels, kw, stride=1, padding=pw),
norm="",
use_spectral_norm=use_spectral_norm))
if out_activation == "":
pass
elif out_activation == "sigmoid":
self.layers.append(nn.Sigmoid())
elif out_activation == "tanh":
self.layers.append(nn.Tanh())
else:
raise ValueError
self.layers = nn.Sequential(*self.layers)
def __init__(self, embed_dim, nf, res_channels, out_channels, norm,
use_spectral_norm, n_layers, out_activation):
super().__init__()
self.upsample_t = norm_layer(
nn.ConvTranspose2d(embed_dim, embed_dim, 4, 2, 1), norm,
use_spectral_norm)
self.dec = ResDecoder(embed_dim + embed_dim,
nf,
res_channels,
out_channels,
norm,
use_spectral_norm,
n_layers,
out_activation,
stride=4)