def __init__(self, in_channels, hidden_channels, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed):
super().__init__()
self.flow_coupling = flow_coupling
self.actnorm = ActNorm2d(in_channels, actnorm_scale)
# 2. permute
if flow_permutation == "invconv":
self.invconv = InvertibleConv1x1(in_channels,
LU_decomposed=LU_decomposed)
self.flow_permutation = \
lambda z, logdet, rev: self.invconv(z, logdet, rev)
elif flow_permutation == "shuffle":
self.shuffle = Permute2d(in_channels, shuffle=True)
self.flow_permutation = \
lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
else:
self.reverse = Permute2d(in_channels, shuffle=False)
self.flow_permutation = \
lambda z, logdet, rev: (self.reverse(z, rev), logdet)
# 3. coupling
if flow_coupling == "additive":
self.block = get_block(in_channels // 2, in_channels // 2,
hidden_channels)
elif flow_coupling == "affine":
self.block = get_block(in_channels // 2, in_channels,
hidden_channels)
def __init__(self, in_channels, hidden_channels, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed,
flow_embed_dim, level):
super().__init__()
self.flow_coupling = flow_coupling
self.actnorm = ActNorm2d(in_channels, actnorm_scale)
self.actnorm_embed = ActNorm2d(flow_embed_dim * (2**level),
actnorm_scale)
if flow_coupling == "additive":
self.conv_proj = Conv2dZeros(hidden_channels,
in_channels // 2,
kernel_size=(3, 3))
elif flow_coupling == "affine":
self.conv_proj = Conv2dZeros(hidden_channels,
in_channels,
kernel_size=(3, 3))
# 2. permute
if flow_permutation == "invconv":
self.invconv = InvertibleConv1x1(in_channels,
LU_decomposed=LU_decomposed)
self.flow_permutation = \
lambda z, logdet, rev: self.invconv(z, logdet, rev)
elif flow_permutation == "shuffle":
self.shuffle = Permute2d(in_channels, shuffle=True)
self.flow_permutation = \
lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
else:
self.reverse = Permute2d(in_channels, shuffle=False)
self.flow_permutation = \
lambda z, logdet, rev: (self.reverse(z, rev), logdet)
self.multgate = nn.Parameter(torch.zeros((6, hidden_channels, 1, 1)))
def __init__(self,
in_channels,
hidden_channels,
actnorm_scale,
flow_permutation,
flow_coupling,
LU_decomposed,
sn,
affine_eps,
no_actnorm,
affine_scale_eps=0,
actnorm_max_scale=None,
no_conv_actnorm=False,
affine_max_scale=0,
actnorm_eps=0):
super().__init__()
self.flow_coupling = flow_coupling
self.affine_eps = affine_eps
self.max_scale = affine_max_scale
assert not (flow_coupling == 'gaffine' and affine_max_scale == 0)
self.no_actnorm = no_actnorm
self.affine_scale_eps = affine_scale_eps
if not self.no_actnorm:
self.actnorm = ActNorm2d(in_channels,
actnorm_scale,
max_scale=actnorm_max_scale,
actnorm_eps=actnorm_eps)
# 2. permute
if flow_permutation == "invconv":
self.invconv = InvertibleConv1x1(in_channels,
LU_decomposed=LU_decomposed)
self.flow_permutation = \
lambda z, logdet, rev: self.invconv(z, logdet, rev)
elif flow_permutation == "shuffle":
self.shuffle = Permute2d(in_channels, shuffle=True)
self.flow_permutation = \
lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
else:
self.reverse = Permute2d(in_channels, shuffle=False)
# 3. coupling
if flow_coupling == "additive":
self.block = get_block(in_channels // 2, in_channels // 2,
hidden_channels, sn, no_conv_actnorm)
elif flow_coupling in ["affine", "naffine", "gaffine"]:
self.block = get_block(in_channels // 2, in_channels,
hidden_channels, sn, no_conv_actnorm)
def __init__(self, in_channels, hidden_channels, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed,
flow_embed_dim, level):
super().__init__()
self.flow_coupling = flow_coupling
self.actnorm = ActNorm2d(in_channels, actnorm_scale)
# 2. permute
if flow_permutation == "invconv":
self.invconv = InvertibleConv1x1(in_channels,
LU_decomposed=LU_decomposed)
self.flow_permutation = \
lambda z, logdet, rev: self.invconv(z, logdet, rev)
elif flow_permutation == "shuffle":
self.shuffle = Permute2d(in_channels, shuffle=True)
self.flow_permutation = \
lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
else:
self.reverse = Permute2d(in_channels, shuffle=False)
self.flow_permutation = \
lambda z, logdet, rev: (self.reverse(z, rev), logdet)