def _create_model(self, n_inputs, n_outputs, rng):
"""
Given input and output sizes, creates and returns the model for the NDE experiments.
"""
model_desc = self.exp_desc.inf.model
if isinstance(model_desc, ed.MDN_Descriptor):
import ml.models.mdns as mdns
return mdns.MDN(n_inputs=n_inputs,
n_outputs=n_outputs,
n_hiddens=model_desc.n_hiddens,
act_fun=model_desc.act_fun,
n_components=model_desc.n_comps,
rng=rng)
elif isinstance(model_desc, ed.MAF_Descriptor):
import ml.models.mafs as mafs
return mafs.ConditionalMaskedAutoregressiveFlow(
n_inputs=n_inputs,
n_outputs=n_outputs,
n_hiddens=model_desc.n_hiddens,
act_fun=model_desc.act_fun,
n_mades=model_desc.n_comps,
mode='random',
rng=rng)
else:
raise TypeError('unknown model descriptor')
def train_maf_cond(n_hiddens, act_fun, n_mades, mode):
assert is_data_loaded(), 'Dataset hasn\'t been loaded'
model = mafs.ConditionalMaskedAutoregressiveFlow(data.n_labels,
data.n_dims,
n_hiddens,
act_fun,
n_mades,
mode=mode)
train_cond(model, a_flow)
save_model(model, 'maf_cond', mode, n_hiddens, act_fun, n_mades, True)