def __init__(self, image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed):
super().__init__()
self.layers = nn.ModuleList()
self.output_shapes = []
self.K = K
self.L = L
H, W, C = image_shape
for i in range(L):
# 1. Squeeze
C, H, W = C * 4, H // 2, W // 2
self.layers.append(SqueezeLayer(factor=2))
self.output_shapes.append([-1, C, H, W])
# 2. K FlowStep
for _ in range(K):
self.layers.append(
FlowStep(in_channels=C,
hidden_channels=hidden_channels,
actnorm_scale=actnorm_scale,
flow_permutation=flow_permutation,
flow_coupling=flow_coupling,
LU_decomposed=LU_decomposed))
self.output_shapes.append([-1, C, H, W])
# 3. Split2d
if i < L - 1:
self.layers.append(Split2d(num_channels=C))
self.output_shapes.append([-1, C // 2, H, W])
C = C // 2
def __init__(self, image_shape, hidden_channels, K, L,
actnorm_scale, flow_permutation, flow_coupling,
LU_decomposed):
super().__init__()
self.layers = nn.ModuleList()
self.output_shapes = []
self.K = K
self.L = L
H, W, C = image_shape
self.estimator = nn.ModuleList()
self.estimator_idx = []
self.flow_idx = []
self.flow_embed_dim = None
flow_counter = 0
for i in range(L):
# 1. Squeeze
C, H, W = C * 4, H // 2, W // 2
self.layers.append(SqueezeLayer(factor=2))
self.output_shapes.append([-1, C, H, W])
self.estimator_idx.append(i)
self.flow_idx.append(flow_counter)
flow_counter += 1
# 3. coupling
if flow_coupling == "additive":
block = Block(C // 2,
C // 2,
hidden_channels)
elif flow_coupling == "affine":
block = Block(C // 2,
C,
hidden_channels)
self.estimator.append(block)
# 2. K FlowStep
for _ in range(K):
self.layers.append(
FlowStep(in_channels=C,
hidden_channels=hidden_channels,
actnorm_scale=actnorm_scale,
flow_permutation=flow_permutation,
flow_coupling=flow_coupling,
LU_decomposed=LU_decomposed,))
self.output_shapes.append([-1, C, H, W])
self.estimator_idx.append(i)
self.flow_idx.append(flow_counter)
flow_counter += 1
# 3. Split2d
if i < L - 1:
self.layers.append(Split2d(num_channels=C))
self.output_shapes.append([-1, C // 2, H, W])
self.estimator_idx.append(i)
self.flow_idx.append(flow_counter)
flow_counter += 1
C = C // 2
def __init__(self, image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed,
logittransform, sn, affine_eps, no_actnorm, affine_scale_eps,
actnorm_max_scale, no_conv_actnorm, affine_max_scale,
actnorm_eps, no_split):
super().__init__()
self.layers = nn.ModuleList()
if logittransform:
self.layers.append(LogitTransform(1e-6))
self.output_shapes = []
self.K = K
self.L = L
H, W, C = image_shape
self.splits = []
for i in range(L):
# 1. Squeeze
C, H, W = C * 4, H // 2, W // 2
self.layers.append(SqueezeLayer(factor=2))
self.output_shapes.append([-1, C, H, W])
# 2. K FlowStep
for _ in range(K):
self.layers.append(
FlowStep(in_channels=C,
hidden_channels=hidden_channels,
actnorm_scale=actnorm_scale,
flow_permutation=flow_permutation,
flow_coupling=flow_coupling,
LU_decomposed=LU_decomposed,
sn=sn,
affine_eps=affine_eps,
no_actnorm=no_actnorm,
affine_scale_eps=affine_scale_eps,
actnorm_max_scale=actnorm_max_scale,
no_conv_actnorm=no_conv_actnorm,
affine_max_scale=affine_max_scale,
actnorm_eps=actnorm_eps))
self.output_shapes.append([-1, C, H, W])
# 3. Split2d
if i < L - 1:
if no_split:
self.output_shapes.append([-1, C, H, W])
else:
split = Split2d(num_channels=C)
self.splits.append(split)
self.layers.append(split)
self.output_shapes.append([-1, C // 2, H, W])
C = C // 2
def __init__(self, image_shape, hidden_channels, K, L,
actnorm_scale, flow_permutation, flow_coupling,
LU_decomposed, extra_condition, sp_condition, num_classes):
super().__init__()
self.layers = nn.ModuleList()
self.output_shapes = []
self.K = K
self.L = L
H, W, C = image_shape
print("image shape", H, W, C)
for i in range(L):
# 1. Squeeze
C, H, W = C * 4, H // 2, W // 2
self.layers.append(SqueezeLayer(factor=2))
self.output_shapes.append([-1, C, H, W])
print("after squeeze", C)
# 2. K FlowStep
for _ in range(K):
self.layers.append(
FlowStep(in_channels=C,
hidden_channels=hidden_channels,
actnorm_scale=actnorm_scale,
flow_permutation=flow_permutation,
flow_coupling=flow_coupling,
LU_decomposed=LU_decomposed,
extra_condition=extra_condition,
num_classes=num_classes))
self.output_shapes.append([-1, C, H, W])
# 3. Split2d
if i < L - 1:
print("split2d num channels", C)
self.layers.append(Split2d(num_channels=C, y_classes=num_classes, sp_condition=sp_condition))
self.output_shapes.append([-1, C // 2, H, W])
C = C // 2
print("after split", C)
