def __init__(self, in_channels=1, mid_channels=64, num_blocks=4):
super(RealNVPMNIST, self).__init__()
self.body = iSequential(
addZslot(), passThrough(iLogits()),
passThrough(
CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(
CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True))),
passThrough(
CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(SqueezeLayer(2)),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=False))),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=True))),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=False))),
keepChannels(2 * in_channels),
passThrough(
CouplingLayer(2 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(
CouplingLayer(2 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True))),
passThrough(
CouplingLayer(2 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(SqueezeLayer(2)),
passThrough(
CouplingLayer(8 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=False))),
passThrough(
CouplingLayer(8 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=True))),
passThrough(
CouplingLayer(8 * in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=False))),
keepChannels(4 * in_channels),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True))),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False))),
passThrough(
CouplingLayer(4 * in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True))),
FlatJoin())
def __init__(self, num_scales=2, in_channels=3, mid_channels=64, num_blocks=8, init_zeros=False,
st_type='resnet', use_batch_norm=True, img_shape=(1, 28, 28), skip=True, latent_dim=None):
super(RealNVPSmall, self).__init__()
self.body = iSequential(
addZslot(),
passThrough(iLogits()),
passThrough(CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=False), init_zeros=init_zeros, st_type=st_type, use_batch_norm=use_batch_norm)),
passThrough(CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True), init_zeros=init_zeros, st_type=st_type, use_batch_norm=use_batch_norm)),
FlatJoin()
)
def __init__(self, num_scales=2, in_channels=3, mid_channels=64, num_blocks=8, init_zeros=False,
st_type='resnet', use_batch_norm=True, img_shape=(1, 28, 28), skip=True, latent_dim=100):
super(RealNVP, self).__init__()
layers = [addZslot(), passThrough(iLogits())]
self.output_shapes = []
_, _, img_width = img_shape
for scale in range(num_scales):
in_couplings = self._threecouplinglayers(in_channels, mid_channels, num_blocks, MaskCheckerboard,
init_zeros, st_type, use_batch_norm, img_width, skip, latent_dim)
layers.append(passThrough(*in_couplings))
if scale == num_scales - 1:
layers.append(passThrough(
CouplingLayer(in_channels, mid_channels, num_blocks, MaskCheckerboard(reverse_mask=True),
init_zeros, st_type, use_batch_norm, img_width, skip, latent_dim)))
else:
layers.append(passThrough(SqueezeLayer(2)))
img_width = img_width // 2
if st_type != 'autoencoder': # in the autoencoder case we probably want the bottleneck size to be fixed?
mid_channels *= 2
out_couplings = self._threecouplinglayers(4*in_channels, mid_channels, num_blocks,
MaskChannelwise, init_zeros, st_type, use_batch_norm, img_width, skip, latent_dim)
layers.append(passThrough(*out_couplings))
layers.append(keepChannels(2*in_channels))
in_channels *= 2
layers.append(FlatJoin())
self.body = iSequential(*layers)
def __init__(self,
num_scales=2,
in_channels=3,
mid_channels=64,
num_blocks=8):
super(RealNVP, self).__init__()
layers = [addZslot(), passThrough(iLogits())]
for scale in range(num_scales):
in_couplings = self._threecouplinglayers(in_channels, mid_channels,
num_blocks,
MaskCheckerboard)
layers.append(passThrough(*in_couplings))
if scale == num_scales - 1:
layers.append(
passThrough(
CouplingLayer(in_channels, mid_channels, num_blocks,
MaskCheckerboard(reverse_mask=True))))
else:
layers.append(passThrough(SqueezeLayer(2)))
out_couplings = self._threecouplinglayers(
4 * in_channels, 2 * mid_channels, num_blocks,
MaskChannelwise)
layers.append(passThrough(*out_couplings))
layers.append(keepChannels(2 * in_channels))
in_channels *= 2
mid_channels *= 2
layers.append(FlatJoin())
self.body = iSequential(*layers)