def __init__(self, data_shape, hidden_size, latent_size, num_res_block,
num_mode, controller_rate):
super().__init__()
self.encoded_shape = (hidden_size[-1],
data_shape[1] // (2**len(hidden_size)),
data_shape[2] // (2**len(hidden_size)))
self.linear = nn.Sequential(
MultimodalController(latent_size, num_mode, controller_rate),
Wrapper(nn.Linear(latent_size,
np.prod(self.encoded_shape).item())),
Wrapper(nn.BatchNorm1d(np.prod(self.encoded_shape).item())),
Wrapper(nn.ReLU(inplace=True)),
)
blocks = [
MultimodalController(hidden_size[-1], num_mode, controller_rate)
]
for i in range(num_res_block):
blocks.append(ResBlock(hidden_size[-1], num_mode, controller_rate))
for i in range(len(hidden_size) - 1, 0, -1):
blocks.extend([
Wrapper(
nn.ConvTranspose2d(hidden_size[i], hidden_size[i - 1], 4,
2, 1)),
Wrapper(nn.BatchNorm2d(hidden_size[i - 1])),
Wrapper(nn.ReLU(inplace=True)),
MultimodalController(hidden_size[i - 1], num_mode,
controller_rate)
])
blocks.extend([
Wrapper(nn.ConvTranspose2d(hidden_size[0], data_shape[0], 4, 2,
1)),
Wrapper(nn.Sigmoid())
])
self.blocks = nn.Sequential(*blocks)
def __init__(self, input_size, output_size, num_mode, controller_rate,
stride):
super().__init__()
self.mc_1 = MultimodalController(input_size, num_mode, controller_rate)
self.mc_2 = MultimodalController(output_size, num_mode,
controller_rate)
self.conv = nn.Sequential(
Wrapper(nn.BatchNorm2d(input_size)),
Wrapper(nn.ReLU()),
Wrapper(nn.Upsample(scale_factor=stride, mode='nearest')),
self.mc_1,
Wrapper(nn.Conv2d(input_size, output_size, 3, 1, 1)),
Wrapper(nn.BatchNorm2d(output_size)),
Wrapper(nn.ReLU()),
self.mc_2,
Wrapper(nn.Conv2d(output_size, output_size, 3, 1, 1)),
)
if stride > 1:
self.shortcut = nn.Sequential(
Wrapper(nn.Upsample(scale_factor=stride,
mode='nearest')), self.mc_1,
Wrapper(nn.Conv2d(input_size, output_size, 1, 1, 0)))
elif input_size != output_size:
self.shortcut = nn.Sequential(
self.mc_1, nn.Conv2d(input_size, output_size, 1, 1, 0))
else:
self.shortcut = nn.Sequential()
def __init__(self, hidden_size, num_mode, controller_rate):
super().__init__()
self.conv = nn.Sequential(
Wrapper(nn.Conv2d(hidden_size, hidden_size, 3, 1, 1)),
Wrapper(nn.BatchNorm2d(hidden_size)),
Wrapper(nn.ReLU(inplace=True)),
MultimodalController(hidden_size, num_mode, controller_rate),
Wrapper(nn.Conv2d(hidden_size, hidden_size, 3, 1, 1)),
Wrapper(nn.BatchNorm2d(hidden_size)),
MultimodalController(hidden_size, num_mode, controller_rate),
)
self.activation = Wrapper(nn.ReLU(inplace=True))
def __init__(self, data_shape, hidden_size, num_mode, controller_rate):
super().__init__()
self.data_shape = data_shape
blocks = [
FirstDisResBlock(data_shape[0], hidden_size[0], num_mode,
controller_rate)
]
if cfg['data_name'] in ['CIFAR10', 'CIFAR100']:
for i in range(len(hidden_size) - 3):
blocks.append(
DisResBlock(hidden_size[i],
hidden_size[i + 1],
num_mode,
controller_rate,
stride=2))
blocks.extend([
DisResBlock(hidden_size[-3],
hidden_size[-2],
num_mode,
controller_rate,
stride=1),
DisResBlock(hidden_size[-2],
hidden_size[-1],
num_mode,
controller_rate,
stride=1),
Wrapper(nn.ReLU()),
MultimodalController(hidden_size[-1], num_mode,
controller_rate),
Wrapper(GlobalSumPooling()),
Wrapper(nn.Linear(hidden_size[-1], 1))
])
else:
for i in range(len(hidden_size) - 2):
blocks.append(
DisResBlock(hidden_size[i],
hidden_size[i + 1],
num_mode,
controller_rate,
stride=2))
blocks.extend([
DisResBlock(hidden_size[-2],
hidden_size[-1],
num_mode,
controller_rate,
stride=1),
Wrapper(nn.ReLU()),
MultimodalController(hidden_size[-1], num_mode,
controller_rate),
Wrapper(GlobalSumPooling()),
Wrapper(nn.Linear(hidden_size[-1], 1))
])
self.blocks = nn.Sequential(*blocks)
def __init__(self, input_size, hidden_size=512, affine=True, num_mode=None, controller_rate=None):
super().__init__()
self.affine = affine
self.net = nn.Sequential(
Wrapper(nn.Conv2d(input_size // 2, hidden_size, 3, padding=1)),
Wrapper(ActNorm(hidden_size, logdet=False)),
Wrapper(nn.ReLU(inplace=True)),
MultimodalController(hidden_size, num_mode, controller_rate),
Wrapper(nn.Conv2d(hidden_size, hidden_size, 1)),
Wrapper(ActNorm(hidden_size, logdet=False)),
Wrapper(nn.ReLU(inplace=True)),
MultimodalController(hidden_size, num_mode, controller_rate),
Wrapper(ZeroConv2d(hidden_size, input_size if self.affine else input_size // 2)),
)
self.net[0].module.weight.data.normal_(0, 0.05)
self.net[0].module.bias.data.zero_()
self.net[4].module.weight.data.normal_(0, 0.05)
self.net[4].module.bias.data.zero_()
def __init__(self, input_size, output_size, num_mode, controller_rate):
super().__init__()
self.mc_1 = MultimodalController(output_size, num_mode,
controller_rate)
self.conv = nn.Sequential(
Wrapper(nn.Conv2d(input_size, output_size, 3, 1, 1)),
Wrapper(nn.ReLU()),
self.mc_1,
Wrapper(nn.Conv2d(output_size, output_size, 3, 1, 1)),
Wrapper(nn.AvgPool2d(2)),
)
self.shortcut = nn.Sequential(
Wrapper(nn.Conv2d(input_size, output_size, 1, 1, 0)),
Wrapper(nn.AvgPool2d(2)),
)
def __init__(self, data_shape, latent_size, hidden_size, num_mode,
controller_rate):
super().__init__()
self.latent_size = latent_size
self.linear = Wrapper(nn.Linear(latent_size, hidden_size[0] * 4 * 4))
blocks = []
for i in range(len(hidden_size) - 1):
blocks.append(
GenResBlock(hidden_size[i],
hidden_size[i + 1],
num_mode,
controller_rate,
stride=2))
blocks.extend([
Wrapper(nn.BatchNorm2d(hidden_size[-1])),
Wrapper(nn.ReLU()),
MultimodalController(hidden_size[-1], num_mode, controller_rate),
Wrapper(nn.Conv2d(hidden_size[-1], data_shape[0], 3, 1, 1)),
Wrapper(nn.Tanh())
])
self.blocks = nn.Sequential(*blocks)