def __init__(self,
in_channels,
hidden_channels=512,
h_channels=0,
inverse=False,
transform='affine',
alpha=1.0,
coupling_type='conv',
h_type=None,
activation='relu',
normalize=None,
num_groups=None,
**kwargs):
super(GlowStep, self).__init__(inverse)
self.actnorm = ActNorm2dFlow(in_channels, inverse=inverse)
self.conv1x1 = Conv1x1Flow(in_channels, inverse=inverse)
self.unit = GlowUnit(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
inverse=inverse,
transform=transform,
alpha=alpha,
coupling_type=coupling_type,
h_type=h_type,
activation=activation,
normalize=normalize,
num_groups=num_groups)
def __init__(self,
in_channels,
kernel_size,
h_channels=0,
inverse=False,
transform='affine',
alpha=1.0,
h_type=None,
activation='relu'):
super(MaCowUnit, self).__init__(inverse)
self.conv1 = MaskedConvFlow(in_channels,
(kernel_size[0], kernel_size[1]),
order='A',
h_channels=h_channels,
transform=transform,
alpha=alpha,
h_type=h_type,
activation=activation)
self.conv2 = MaskedConvFlow(in_channels,
(kernel_size[0], kernel_size[1]),
order='B',
h_channels=h_channels,
transform=transform,
alpha=alpha,
h_type=h_type,
activation=activation)
self.actnorm1 = ActNorm2dFlow(in_channels, inverse=inverse)
self.conv3 = MaskedConvFlow(in_channels,
(kernel_size[1], kernel_size[0]),
order='C',
h_channels=h_channels,
transform=transform,
alpha=alpha,
h_type=h_type,
activation=activation)
self.conv4 = MaskedConvFlow(in_channels,
(kernel_size[1], kernel_size[0]),
order='D',
h_channels=h_channels,
transform=transform,
alpha=alpha,
h_type=h_type,
activation=activation)
self.actnorm2 = ActNorm2dFlow(in_channels, inverse=inverse)
def __init__(self, in_channels, hidden_channels, h_channels, factor, transform, alpha,
inverse, coupling_type, h_type, activation, normalize, num_groups):
super(MultiScalePrior, self).__init__(inverse)
self.conv1x1 = Conv1x1Flow(in_channels, inverse=inverse)
self.coupling = NICE2d(in_channels, hidden_channels=hidden_channels, h_channels=h_channels,
transform=transform, alpha=alpha, inverse=inverse, factor=factor,
type=coupling_type, h_type=h_type, split_type='continuous', order='up',
activation=activation, normalize=normalize, num_groups=num_groups)
out_channels = in_channels // factor
self.z1_channels = self.coupling.z1_channels
assert out_channels + self.z1_channels == in_channels
self.actnorm = ActNorm2dFlow(out_channels, inverse=inverse)
def __init__(self, in_channels, hidden_channels=512, h_channels=0, inverse=False,
transform='affine', alpha=1.0, coupling_type='conv', h_type=None,
activation='relu', normalize=None, num_groups=None, mask_pos=None):
# mask_pos can be [0, 1, None]
if mask_pos:
assert mask_pos in [0, 1], 'mask position can only be 0 or 1'
super(GlowUnit, self).__init__(inverse)
self.coupling1_up = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='continuous', order='up',
activation=activation, normalize=normalize, num_groups=num_groups)
self.coupling1_dn = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='continuous', order='down',
activation=activation, normalize=normalize, num_groups=num_groups)
self.actnorm = ActNorm2dFlow(in_channels, inverse=inverse)
self.coupling2_up = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='skip', order='up',
activation=activation, normalize=normalize, num_groups=num_groups)
self.coupling2_dn = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='skip', order='down',
activation=activation, normalize=normalize, num_groups=num_groups)
h_type = 'global_mask'
if mask_pos == 0:
self.coupling1_up = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='continuous', order='up',
activation=activation, normalize=normalize, num_groups=num_groups)
self.coupling1_dn = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='continuous', order='down',
activation=activation, normalize=normalize, num_groups=num_groups)
if mask_pos == 1:
self.coupling2_up = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='skip', order='up',
activation=activation, normalize=normalize, num_groups=num_groups)
self.coupling2_dn = NICE2d(in_channels, hidden_channels=hidden_channels,
h_channels=h_channels, transform=transform, alpha=alpha, inverse=inverse,
type=coupling_type, h_type=h_type, split_type='skip', order='down',
activation=activation, normalize=normalize, num_groups=num_groups)
def __init__(self, in_channels, hidden_channels=512, h_channels=0, inverse=False,
transform='affine', alpha=1.0, coupling_type='conv', h_type=None,
activation='relu', normalize=None, num_groups=None, mask_pos=None, conv1x1_type='conv1x1', **kwargs):
super(GlowStep, self).__init__(inverse)
self.actnorm = ActNorm2dFlow(in_channels, inverse=inverse)
if conv1x1_type == 'conv1x1':
self.conv1x1 = Conv1x1Flow(in_channels, inverse=inverse)
elif conv1x1_type == 'cond_conv1x1':
self.conv1x1 = CondConv1x1Flow(in_channels, h_channels, inverse=inverse)
else:
raise ValueError('Please give valid conv1x1 type')
self.conv1x1_type = conv1x1_type
self.unit = GlowUnit(in_channels, hidden_channels=hidden_channels, h_channels=h_channels,
inverse=inverse, transform=transform, alpha=alpha, coupling_type=coupling_type,
h_type=h_type, activation=activation, normalize=normalize, num_groups=num_groups, mask_pos=mask_pos)
def __init__(self,
in_channels,
hidden_channels=512,
h_channels=0,
inverse=False,
transform='affine',
alpha=1.0,
coupling_type='conv',
h_type=None,
activation='relu',
normalize=None,
num_groups=None):
super(GlowUnit, self).__init__(inverse)
self.coupling1_up = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='continuous',
order='up',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.coupling1_dn = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='continuous',
order='down',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.actnorm = ActNorm2dFlow(in_channels, inverse=inverse)
self.coupling2_up = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='skip',
order='up',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.coupling2_dn = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='skip',
order='down',
activation=activation,
normalize=normalize,
num_groups=num_groups)
def __init__(self,
in_channels,
kernel_size,
hidden_channels,
h_channels,
inverse=False,
transform='affine',
alpha=1.0,
coupling_type='conv',
h_type=None,
activation='relu',
normalize=None,
num_groups=None,
**kwargs):
super(MaCowStep, self).__init__(inverse)
num_units = 2
self.actnorm1 = ActNorm2dFlow(in_channels, inverse=inverse)
self.conv1x1 = Conv1x1Flow(in_channels, inverse=inverse)
units = [
MaCowUnit(in_channels,
kernel_size,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
h_type=h_type,
activation=activation) for _ in range(num_units)
]
self.units1 = nn.ModuleList(units)
self.coupling1_up = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='continuous',
order='up',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.coupling1_dn = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='continuous',
order='down',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.actnorm2 = ActNorm2dFlow(in_channels, inverse=inverse)
units = [
MaCowUnit(in_channels,
kernel_size,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
h_type=h_type,
activation=activation) for _ in range(num_units)
]
self.units2 = nn.ModuleList(units)
self.coupling2_up = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='skip',
order='up',
activation=activation,
normalize=normalize,
num_groups=num_groups)
self.coupling2_dn = NICE2d(in_channels,
hidden_channels=hidden_channels,
h_channels=h_channels,
transform=transform,
alpha=alpha,
inverse=inverse,
type=coupling_type,
h_type=h_type,
split_type='skip',
order='down',
activation=activation,
normalize=normalize,
num_groups=num_groups)