def _threecouplinglayers(in_channels, mid_channels, num_blocks,
mask_class):
layers = [
iConv1x1(in_channels),
CouplingLayer(in_channels, mid_channels, num_blocks,
mask_class(reverse_mask=False)),
iConv1x1(in_channels),
CouplingLayer(in_channels, mid_channels, num_blocks,
mask_class(reverse_mask=True)),
iConv1x1(in_channels),
CouplingLayer(in_channels, mid_channels, num_blocks,
mask_class(reverse_mask=False))
]
return layers
def _glow_step(in_channels, mid_channels, num_blocks):
layers = [
ActNorm(in_channels),
iConv1x1(in_channels),
CouplingLayer(in_channels, mid_channels, num_blocks,
MaskChannelwise(reverse_mask=False)),
]
return layers
def _threecouplinglayers(in_channels,
mid_channels,
num_blocks,
mask_class,
num_classes=10):
layers = []
for i in range(3):
layers.append(
passThrough(
iConv1x1(in_channels),
CouplingLayer(in_channels, mid_channels, num_blocks,
mask_class(reverse_mask=not i % 2)),
ActNorm2d(in_channels)))
layers.append(iCategoricalFiLM(num_classes, in_channels))
return layers