def __init__(self,
in_channels,
out_channels,
kernel_size=(3, 3),
stride=(1, 1),
padding="same",
do_actnorm=True,
weight_std=0.05):
super().__init__()
if padding == "same":
padding = compute_same_pad(kernel_size, stride)
elif padding == "valid":
padding = 0
self.conv = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding,
bias=(not do_actnorm))
# init weight with std
self.conv.weight.data.normal_(mean=0.0, std=weight_std)
if not do_actnorm:
self.conv.bias.data.zero_()
else:
self.actnorm = ActNorm(out_channels, scale_fn='exp', eps=0)
self.do_actnorm = do_actnorm
def flow_block_down_1(n_chans, hidden_channels):
in_to_outs = [[
nn.Sequential(
nn.Conv1d(n_chans * 2, hidden_channels, 5, padding=5 // 2)),
nn.Sequential(
nn.Conv1d(n_chans * 8, hidden_channels, 3, padding=3 // 2),
nn.Upsample(scale_factor=2))
]]
m = nn.Sequential(
MultipleInputOutput(in_to_outs),
Expression(unwrap_single_element),
nn.Conv1d(hidden_channels, n_chans * 4 // 2, 5, padding=5 // 2),
MultiplyFactors(n_chans * 4 // 2),
)
c = CouplingLayer(ChunkChansIn2(swap_dims=False),
AdditiveCoefs(m),
AffineModifier('sigmoid', add_first=True, eps=0),
condition_merger=CatCondAsList(cond_preproc=None))
f = InvertibleSequential(
ActNorm(
n_chans * 4,
'exp',
),
InvPermute(n_chans * 4, fixed=False, use_lu=True),
c,
)
return f
def dense_flow_block(n_chans, hidden_channels):
return InvertibleSequential(
ActNorm(
n_chans,
'exp',
), InvPermute(n_chans, fixed=False, use_lu=True),
CouplingLayer(
ChunkChansIn2(swap_dims=False),
AdditiveCoefs(
nn.Sequential(
nn.Linear(n_chans // 2, hidden_channels),
nn.ELU(),
nn.Linear(hidden_channels, n_chans // 2),
MultiplyFactors(n_chans // 2),
)), AffineModifier('sigmoid', add_first=True, eps=0)))
def conv_flow_block(n_chans, hidden_channels, kernel_length):
assert kernel_length % 2 == 1
return InvertibleSequential(
ActNorm(
n_chans,
'exp',
), InvPermute(n_chans, fixed=False, use_lu=True),
CouplingLayer(
ChunkChansIn2(swap_dims=False),
AdditiveCoefs(
nn.Sequential(
nn.Conv1d(n_chans // 2,
hidden_channels,
kernel_length,
padding=kernel_length // 2),
nn.ELU(),
nn.Conv1d(hidden_channels,
n_chans // 2,
kernel_length,
padding=kernel_length // 2),
MultiplyFactors(n_chans // 2),
)), AffineModifier('sigmoid', add_first=True, eps=0)))
def create_glow_model(hidden_channels,
K,
L,
flow_permutation,
flow_coupling,
LU_decomposed,
n_chans,
block_type='conv',
use_act_norm=True):
image_shape = (32, 32, n_chans)
H, W, C = image_shape
flows_per_scale = []
act_norms_per_scale = []
dists_per_scale = []
for i in range(L):
C, H, W = C * 4, H // 2, W // 2
splitter = SubsampleSplitter(2,
via_reshape=True,
chunk_chans_first=True,
checkerboard=False,
cat_at_end=True)
if block_type == 'dense':
pre_flow_layers = [Flatten2d()]
in_channels = C * H * W
else:
assert block_type == 'conv'
pre_flow_layers = []
in_channels = C
flow_layers = [
flow_block(in_channels=in_channels,
hidden_channels=hidden_channels,
flow_permutation=flow_permutation,
flow_coupling=flow_coupling,
LU_decomposed=LU_decomposed,
cond_channels=0,
cond_merger=None,
block_type=block_type,
use_act_norm=use_act_norm) for _ in range(K)
]
if block_type == 'dense':
post_flow_layers = [ViewAs((-1, C * H * W), (-1, C, H, W))]
else:
assert block_type == 'conv'
post_flow_layers = []
flow_layers = pre_flow_layers + flow_layers + post_flow_layers
flow_this_scale = InvertibleSequential(splitter, *flow_layers)
flows_per_scale.append(flow_this_scale)
if i < L - 1:
# there will be a chunking here
C = C // 2
# act norms for distribution (mean/std as actnorm isntead of integrated
# into dist)
act_norms_per_scale.append(
InvertibleSequential(Flatten2d(),
ActNorm((C * H * W), scale_fn='exp')))
dists_per_scale.append(
Unlabeled(
NClassIndependentDist(1,
C * H * W,
optimize_mean=False,
optimize_std=False)))
assert len(flows_per_scale) == 3
nd_1_o = Node(None, flows_per_scale[0], name='m0-flow-0')
nd_1_ab = Node(nd_1_o, ChunkChans(2))
nd_1_a = SelectNode(nd_1_ab, 0)
nd_1_an = Node(nd_1_a, act_norms_per_scale[0], name='m0-act-0')
nd_1_ad = Node(nd_1_an, dists_per_scale[0], name='m0-dist-0')
nd_1_b = SelectNode(nd_1_ab, 1, name='m0-in-flow-1')
nd_2_o = Node(nd_1_b, flows_per_scale[1], name='m0-flow-1')
nd_2_ab = Node(
nd_2_o,
ChunkChans(2),
)
nd_2_a = SelectNode(
nd_2_ab,
0,
)
nd_2_an = Node(nd_2_a, act_norms_per_scale[1], name='m0-act-1')
nd_2_ad = Node(nd_2_an, dists_per_scale[1], name='m0-dist-1')
nd_2_b = SelectNode(nd_2_ab, 1, name='m0-in-flow-2')
nd_3_o = Node(nd_2_b, flows_per_scale[2], name='m0-flow-2')
nd_3_n = Node(nd_3_o, act_norms_per_scale[2], name='m0-act-2')
nd_3_d = Node(nd_3_n, dists_per_scale[2], name='m0-dist-2')
model = CatAsListNode([nd_1_ad, nd_2_ad, nd_3_d],
name='m0-full') # cahnged to pre-full
return model
def flow_block(in_channels,
hidden_channels,
flow_permutation,
flow_coupling,
LU_decomposed,
cond_channels,
cond_merger,
block_type,
use_act_norm,
nonlin_name='relu'):
if use_act_norm:
actnorm = ActNorm(in_channels, scale_fn='exp', eps=0)
# 2. permute
if flow_permutation == "invconv":
flow_permutation = InvPermute(in_channels,
fixed=False,
use_lu=LU_decomposed)
elif flow_permutation == 'invconvfixed':
flow_permutation = InvPermute(in_channels,
fixed=True,
use_lu=LU_decomposed)
elif flow_permutation == "identity":
flow_permutation = Identity()
else:
assert flow_permutation == 'shuffle'
flow_permutation = Shuffle(in_channels)
if flow_coupling == "additive":
out_channels = in_channels // 2
else:
out_channels = in_channels
if type(block_type) is str:
if block_type == 'conv':
block_fn = get_conv_block
else:
assert block_type == 'dense'
block_fn = get_dense_block
else:
block_fn = block_type
block = block_fn(in_channels // 2 + cond_channels,
out_channels,
hidden_channels,
nonlin_name=nonlin_name)
if flow_coupling == "additive":
coupling = CouplingLayer(ChunkChansIn2(swap_dims=True),
AdditiveCoefs(block, ),
AffineModifier(
sigmoid_or_exp_scale='sigmoid',
eps=0,
add_first=True,
),
condition_merger=cond_merger)
elif flow_coupling == "affine":
coupling = CouplingLayer(
ChunkChansIn2(swap_dims=True),
AffineCoefs(block, EverySecondChan()),
AffineModifier(sigmoid_or_exp_scale='sigmoid',
eps=0,
add_first=True),
condition_merger=cond_merger,
)
else:
assert False, f"unknown flow_coupling {flow_coupling}"
if use_act_norm:
sequential = InvertibleSequential(actnorm, flow_permutation, coupling)
else:
sequential = InvertibleSequential(flow_permutation, coupling)
return sequential