def load_einet_state(model_path,
einet_file='einet.pth',
graph_file='einet.rg',
n_vars=None, n_classes=None, n_sums=None, n_input_dists=None,
exp_fam=None, exp_fam_args=None,
use_em=None, em_freq=None, em_stepsize=None,
graph=None):
# reload model
einet = None
if graph is None:
if graph_file:
graph_file = os.path.join(model_path, graph_file)
graph = Graph.read_gpickle(graph_file)
else:
raise ValueError(f"Cannot create graph")
model_file = os.path.join(model_path, einet_file)
einet = make_einet(graph,
n_vars=n_vars,
n_classes=n_classes,
n_sums=n_sums,
n_input_dists=n_input_dists,
exp_fam=exp_fam, exp_fam_args=exp_fam_args,
use_em=use_em,
em_freq=em_freq, em_stepsize=em_stepsize)
einet.load_state_dict(torch.load(model_file))
print("Loaded model from {}".format(model_file))
return einet, graph
def load_einet(model_path, einet_file='einet.pth', graph_file='einet.rg'):
# reload model
einet, graph = None, None
model_file = os.path.join(model_path, einet_file)
einet = torch.load(model_file)
print("Loaded model from {}".format(model_file))
if graph_file:
graph_file = os.path.join(model_path, graph_file)
graph = Graph.read_gpickle(graph_file)
return einet, graph
