'n_reflections': c['ndim_x']
},
name=f'perm_{i+1}'))
nodes.append(
Node(nodes[-1],
AffineCoupling, {
'F_class': F_fully_connected,
'F_args': {
'internal_size': c['hidden_layer_sizes']
}
},
name=f'ac_{i+1}'))
nodes.append(OutputNode(nodes[-1], name='z'))
model = ReversibleGraphNet(nodes, verbose=False)
model.to(c['device'])
model.params_trainable = list(
filter(lambda p: p.requires_grad, model.parameters()))
def model_inverse(test_z):
return model([test_z], rev=True)
c['model'] = model
c['model_inverse'] = model_inverse
# create namedtuple from config dictionary
c = namedtuple("Configuration", c.keys())(*c.values())
# create namedtuple from config dictionary
c = namedtuple("Configuration",c.keys())(*c.values())
##############################
### MODEL ARCHITECTURE ###
##############################
nodes = [InputNode(c.ndim_x, name='x')]
for i in range(c.n_blocks):
nodes.append(Node(nodes[-1],
HouseholderPerm,
{'fixed': False, 'n_reflections': c.ndim_x},
name=f'perm_{i+1}'))
nodes.append(Node(nodes[-1],
AffineCoupling,
{'F_class': F_fully_connected,
'F_args': {'internal_size': c.hidden_layer_sizes}},
name=f'ac_{i+1}'))
nodes.append(OutputNode(nodes[-1], name='z'))
model = ReversibleGraphNet(nodes, verbose=False)
model.to(c.device)
def model_inverse(test_z):
return model([test_z], rev=True)
