def train_manifold_flow_sequential(args, dataset, model, simulator):
""" MFMF-A training """
assert not args.specified
trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
common_kwargs = {
"dataset": dataset,
"batch_size": args.batchsize,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip,
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
logger.info("Starting training MF, phase 1: manifold training")
learning_curves = trainer.train(
loss_functions=[losses.mse],
loss_labels=["MSE"],
loss_weights=[args.msefactor],
epochs=args.epochs // 2,
parameters=list(model.outer_transform.parameters()) +
list(model.encoder.parameters())
if args.algorithm == "emf" else model.outer_transform.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_A{}.pt")
],
forward_kwargs={"mode": "projection"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
logger.info("Starting training MF, phase 2: density training")
learning_curves_ = trainer.train(
loss_functions=[losses.nll],
loss_labels=["NLL"],
loss_weights=[args.nllfactor],
epochs=args.epochs - (args.epochs // 2),
parameters=model.inner_transform.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_B{}.pt")
],
forward_kwargs={"mode": "mf-fixed-manifold"},
**common_kwargs,
)
learning_curves_ = np.vstack(learning_curves_).T
learning_curves = learning_curves_ if learning_curves is None else np.vstack(
(learning_curves, learning_curves_))
return learning_curves
def train_dough(args, dataset, model, simulator):
""" PIE with variable epsilons training """
trainer = VariableDimensionManifoldFlowTrainer(
model) if simulator.parameter_dim(
) is None else ConditionalVariableDimensionManifoldFlowTrainer(model)
common_kwargs = {
"dataset": dataset,
"batch_size": args.batchsize,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip,
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
logger.info(
"Starting training dough, phase 1: NLL without latent regularization")
learning_curves = trainer.train(
loss_functions=[losses.nll],
loss_labels=["NLL"],
loss_weights=[args.nllfactor],
epochs=args.epochs,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_{}.pt")
],
l1=args.doughl1reg,
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
return learning_curves
def train_pie(args, dataset, model, simulator):
""" PIE training """
trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
logger.info("Starting training PIE on NLL")
common_kwargs = {
"dataset": dataset,
"batch_size": args.batchsize,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip,
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
learning_curves = trainer.train(
loss_functions=[losses.nll],
loss_labels=["NLL"],
loss_weights=[args.nllfactor],
epochs=args.epochs,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_{}.pt")
],
forward_kwargs={"mode": "pie"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
return learning_curves
def train_manifold_flow_alternating(args, dataset, model, simulator):
""" MFMF-A training """
assert not args.specified
trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
metatrainer = AlternatingTrainer(model, trainer, trainer)
meta_kwargs = {
"dataset": dataset,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR
}
if args.weightdecay is not None:
meta_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
phase1_kwargs = {
"forward_kwargs": {
"mode": "projection"
},
"clip_gradient": args.clip
}
phase2_kwargs = {
"forward_kwargs": {
"mode": "mf-fixed-manifold"
},
"clip_gradient": args.clip
}
phase1_parameters = (list(model.outer_transform.parameters()) +
list(model.encoder.parameters()) if args.algorithm
== "emf" else model.outer_transform.parameters())
phase2_parameters = model.inner_transform.parameters()
logger.info(
"Starting training MF, alternating between reconstruction error and log likelihood"
)
learning_curves_ = metatrainer.train(
loss_functions=[losses.mse, losses.nll],
loss_function_trainers=[0, 1],
loss_labels=["MSE", "NLL"],
loss_weights=[args.msefactor, args.nllfactor],
epochs=args.epochs // 2,
subsets=args.subsets,
batch_sizes=[args.batchsize, args.batchsize],
parameters=[phase1_parameters, phase2_parameters],
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_{}.pt")
],
trainer_kwargs=[phase1_kwargs, phase2_kwargs],
**meta_kwargs,
)
learning_curves = np.vstack(learning_curves_).T
return learning_curves
def train_generative_adversarial_manifold_flow_alternating(
args, dataset, model, simulator):
""" MFMF-OTA training """
assert not args.specified
gen_trainer = GenerativeTrainer(model) if simulator.parameter_dim(
) is None else ConditionalGenerativeTrainer(model)
likelihood_trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
metatrainer = AlternatingTrainer(model, gen_trainer, likelihood_trainer)
meta_kwargs = {
"dataset": dataset,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR
}
if args.weightdecay is not None:
meta_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
phase1_kwargs = {"clip_gradient": args.clip}
phase2_kwargs = {
"forward_kwargs": {
"mode": "mf-fixed-manifold"
},
"clip_gradient": args.clip
}
phase1_parameters = model.parameters()
phase2_parameters = model.inner_transform.parameters()
logger.info(
"Starting training GAMF, alternating between Sinkhorn divergence and log likelihood"
)
learning_curves_ = metatrainer.train(
loss_functions=[losses.make_sinkhorn_divergence(), losses.nll],
loss_function_trainers=[0, 1],
loss_labels=["GED", "NLL"],
loss_weights=[args.sinkhornfactor, args.nllfactor],
batch_sizes=[args.genbatchsize, args.batchsize],
epochs=args.epochs // 2,
parameters=[phase1_parameters, phase2_parameters],
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_{}.pt")
],
trainer_kwargs=[phase1_kwargs, phase2_kwargs],
subsets=args.subsets,
subset_callbacks=[callbacks.print_mf_weight_statistics()]
if args.debug else None,
**meta_kwargs,
)
learning_curves = np.vstack(learning_curves_).T
return learning_curves
def sample_from_model(args, model, simulator):
""" Generate samples from model and store """
logger.info("Sampling from model")
if simulator.parameter_dim() is None:
x_gen = model.sample(n=args.generate).detach().numpy()
else:
params = simulator.default_parameters(true_param_id=args.trueparam)
params = np.asarray([params for _ in range(args.generate)])
params = torch.tensor(params, dtype=torch.float)
x_gen = model.sample(n=args.generate, context=params).detach().numpy()
np.save(create_filename("results", "samples", args), x_gen)
return x_gen
def objective(trial):
global counter
counter += 1
# Hyperparameters
margs = pick_parameters(args, trial, counter)
logger.info("Starting training for the following hyperparameters:")
for k, v in margs.__dict__.items():
logger.info(" %s: %s", k, v)
# Bug fix related to some num_workers > 1 and CUDA. Bad things happen otherwise!
torch.multiprocessing.set_start_method("spawn", force=True)
# Load data
simulator = load_simulator(margs)
dataset = load_training_dataset(simulator, margs)
# Create model
model = create_model(margs, simulator)
# Train
_ = train.train_model(margs, dataset, model, simulator)
# Save
torch.save(model.state_dict(), create_filename("model", None, margs))
# Evaluate
model.eval()
# Evaluate test samples
log_likelihood_test, reconstruction_error_test, _ = evaluate.evaluate_test_samples(
margs, simulator, model, paramscan=True)
mean_log_likelihood_test = np.mean(log_likelihood_test)
mean_reco_error_test = np.mean(reconstruction_error_test)
# Generate samples
x_gen = evaluate.sample_from_model(margs, model, simulator)
distances_gen = simulator.distance_from_manifold(x_gen)
mean_gen_distance = np.mean(distances_gen)
# Report results
logger.info("Results:")
logger.info(" test log p: %s", mean_log_likelihood_test)
logger.info(" test reco err: %s", mean_reco_error_test)
logger.info(" gen distance: %s", mean_gen_distance)
return (-1.0 * margs.metricnllfactor * mean_log_likelihood_test +
margs.metricrecoerrorfactor * mean_reco_error_test +
margs.metricdistancefactor * mean_gen_distance)
def evaluate_model_samples(args, simulator, x_gen):
""" Evaluate model samples and save results """
logger.info("Calculating likelihood of generated samples")
try:
if simulator.parameter_dim() is None:
log_likelihood_gen = simulator.log_density(x_gen)
else:
params = simulator.default_parameters(true_param_id=args.trueparam)
params = np.asarray([params for _ in range(args.generate)])
log_likelihood_gen = simulator.log_density(x_gen, parameters=params)
log_likelihood_gen[np.isnan(log_likelihood_gen)] = -1.0e-12
np.save(create_filename("results", "samples_likelihood", args), log_likelihood_gen)
except IntractableLikelihoodError:
logger.info("True simulator likelihood is intractable for dataset %s", args.dataset)
# Distance from manifold
try:
logger.info("Calculating distance from manifold of generated samples")
distances_gen = simulator.distance_from_manifold(x_gen)
np.save(create_filename("results", "samples_manifold_distance", args), distances_gen)
except NotImplementedError:
logger.info("Cannot calculate distance from manifold for dataset %s", args.dataset)
def timing(args):
logger.info(
"Timing algorithm %s with %s outer layers with transformation %s and %s inner layers with transformation %s",
args.algorithm,
args.outerlayers,
args.outertransform,
args.innerlayers,
args.innertransform,
)
# Bug fix related to some num_workers > 1 and CUDA. Bad things happen otherwise!
torch.multiprocessing.set_start_method("spawn", force=True)
if torch.cuda.is_available():
torch.set_default_tensor_type("torch.cuda.DoubleTensor")
# Loop over data dims
all_times = []
for datadim in args.datadims:
logger.info("Starting timing for %s-dimensional data", datadim)
args.datadim = datadim
# Data
data = torch.randn(args.batchsize, datadim)
data.requires_grad = True
# Model
model = create_model(args, context_features=None)
if torch.cuda.is_available():
model = model.to(torch.device("cuda"))
# Time forward pass
times = []
for _ in range(args.repeats):
time_before = time.time()
_ = model(data)
times.append(time.time() - time_before)
logger.info("Mean time: %s s", np.mean(times))
all_times.append(times)
# Save results
logger.info("Saving results")
np.save(create_filename("timing", None, args), all_times)
def train_generative_adversarial_manifold_flow(args, dataset, model,
simulator):
""" MFMF-OT training """
gen_trainer = GenerativeTrainer(model) if simulator.parameter_dim(
) is None else ConditionalGenerativeTrainer(model)
common_kwargs = {
"dataset": dataset,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
logger.info("Starting training GAMF: Sinkhorn-GAN")
callbacks_ = [
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] + "_epoch_{}.pt")
]
if args.debug:
callbacks_.append(callbacks.print_mf_weight_statistics())
learning_curves_ = gen_trainer.train(
loss_functions=[losses.make_sinkhorn_divergence()],
loss_labels=["GED"],
loss_weights=[args.sinkhornfactor],
epochs=args.epochs,
callbacks=callbacks_,
batch_size=args.genbatchsize,
compute_loss_variance=True,
**common_kwargs,
)
learning_curves = np.vstack(learning_curves_).T
return learning_curves
# Simulator
simulator = load_simulator(args)
# Parameters?
conditional = simulator.parameter_dim() is not None
parameters_train = simulator.sample_from_prior(
args.train) if conditional else None
# Sample
if args.train > 0:
logger.info("Generating %s training samples at parameters %s",
args.train, parameters_train)
x_train = simulator.sample(args.train, parameters=parameters_train)
np.save(create_filename("sample", "x_train", args), x_train)
if conditional:
np.save(create_filename("sample", "parameters_train", args),
parameters_train)
if args.paramscan > 0:
parameters_val = np.array(
[simulator.default_parameters()
for _ in range(args.paramscan)]).reshape(
(args.paramscan, -1)) if conditional else None
logger.info("Generating %s param-scan samples at parameters %s",
args.paramscan, parameters_val)
x_val = simulator.sample(args.paramscan, parameters=parameters_val)
np.save(create_filename("sample", "x_paramscan", args), x_val)
if conditional:
np.save(create_filename("sample", "parameters_paramscan", args),
def train_slice_of_pie(args, dataset, model, simulator):
""" SLICE training """
trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
common_kwargs = {
"dataset": dataset,
"batch_size": args.batchsize,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip,
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
if args.nopretraining or args.epochs // 3 < 1:
logger.info("Skipping pretraining phase")
learning_curves = np.zeros((0, 2))
else:
logger.info(
"Starting training slice of PIE, phase 1: pretraining on reconstruction error"
)
learning_curves = trainer.train(
loss_functions=[losses.mse],
loss_labels=["MSE"],
loss_weights=[args.initialmsefactor],
epochs=args.epochs // 3,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_A{}.pt")
],
forward_kwargs={"mode": "projection"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
logger.info("Starting training slice of PIE, phase 2: mixed training")
learning_curves_ = trainer.train(
loss_functions=[losses.mse, losses.nll],
loss_labels=["MSE", "NLL"],
loss_weights=[args.initialmsefactor, args.initialnllfactor],
epochs=args.epochs - (1 if args.nopretraining else 2) *
(args.epochs // 3),
parameters=model.inner_transform.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_B{}.pt")
],
forward_kwargs={"mode": "slice"},
**common_kwargs,
)
learning_curves_ = np.vstack(learning_curves_).T
learning_curves = np.vstack((learning_curves, learning_curves_))
logger.info(
"Starting training slice of PIE, phase 3: training only inner flow on NLL"
)
learning_curves_ = trainer.train(
loss_functions=[losses.mse, losses.nll],
loss_labels=["MSE", "NLL"],
loss_weights=[args.msefactor, args.nllfactor],
epochs=args.epochs // 3,
parameters=model.inner_transform.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_C{}.pt")
],
forward_kwargs={"mode": "slice"},
**common_kwargs,
)
learning_curves_ = np.vstack(learning_curves_).T
learning_curves = np.vstack((learning_curves, learning_curves_))
return learning_curves
def train_manifold_flow(args, dataset, model, simulator):
""" MFMF-S training """
trainer = ManifoldFlowTrainer(model) if simulator.parameter_dim(
) is None else ConditionalManifoldFlowTrainer(model)
common_kwargs = {
"dataset": dataset,
"batch_size": args.batchsize,
"initial_lr": args.lr,
"scheduler": optim.lr_scheduler.CosineAnnealingLR,
"clip_gradient": args.clip,
}
if args.weightdecay is not None:
common_kwargs["optimizer_kwargs"] = {
"weight_decay": float(args.weightdecay)
}
if args.specified:
logger.info("Starting training MF with specified manifold on NLL")
learning_curves = trainer.train(
loss_functions=[losses.mse, losses.nll],
loss_labels=["MSE", "NLL"],
loss_weights=[0.0, args.nllfactor],
epochs=args.epochs,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_{}.pt")
],
forward_kwargs={"mode": "mf"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
else:
if args.nopretraining or args.epochs // args.prepostfraction < 1:
logger.info("Skipping pretraining phase")
learning_curves = None
elif args.prepie:
logger.info(
"Starting training MF, phase 1: pretraining on PIE likelihood")
learning_curves = trainer.train(
loss_functions=[losses.nll],
loss_labels=["NLL"],
loss_weights=[args.nllfactor],
epochs=args.epochs // args.prepostfraction,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_A{}.pt")
],
forward_kwargs={"mode": "pie"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
else:
logger.info(
"Starting training MF, phase 1: pretraining on reconstruction error"
)
learning_curves = trainer.train(
loss_functions=[losses.mse],
loss_labels=["MSE"],
loss_weights=[args.msefactor],
epochs=args.epochs // args.prepostfraction,
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_A{}.pt")
],
forward_kwargs={"mode": "projection"},
**common_kwargs,
)
learning_curves = np.vstack(learning_curves).T
logger.info("Starting training MF, phase 2: mixed training")
learning_curves_ = trainer.train(
loss_functions=[losses.mse, losses.nll],
loss_labels=["MSE", "NLL"],
loss_weights=[args.msefactor, args.addnllfactor],
epochs=args.epochs -
(2 - int(args.nopretraining) - int(args.noposttraining)) *
(args.epochs // args.prepostfraction),
parameters=model.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_B{}.pt")
],
forward_kwargs={"mode": "mf"},
**common_kwargs,
)
learning_curves_ = np.vstack(learning_curves_).T
learning_curves = learning_curves_ if learning_curves is None else np.vstack(
(learning_curves, learning_curves_))
if args.nopretraining or args.epochs // args.prepostfraction < 1:
logger.info("Skipping inner flow phase")
else:
logger.info(
"Starting training MF, phase 3: training only inner flow on NLL"
)
learning_curves_ = trainer.train(
loss_functions=[losses.mse, losses.nll],
loss_labels=["MSE", "NLL"],
loss_weights=[0.0, args.nllfactor],
epochs=args.epochs // args.prepostfraction,
parameters=model.inner_transform.parameters(),
callbacks=[
callbacks.save_model_after_every_epoch(
create_filename("checkpoint", None, args)[:-3] +
"_epoch_C{}.pt")
],
forward_kwargs={"mode": "mf-fixed-manifold"},
**common_kwargs,
)
learning_curves_ = np.vstack(learning_curves_).T
learning_curves = np.vstack((learning_curves, learning_curves_))
return learning_curves
# Bug fix related to some num_workers > 1 and CUDA. Bad things happen otherwise!
torch.multiprocessing.set_start_method("spawn", force=True)
# Data
simulator = load_simulator(args)
dataset = load_training_dataset(simulator, args)
# Model
model = create_model(args, simulator)
# Maybe load pretrained model
if args.load is not None:
args_ = copy.deepcopy(args)
args_.modelname = args.load
if args_.i > 0:
args_.modelname += "_run{}".format(args_.i)
logger.info("Loading model %s", args_.modelname)
model.load_state_dict(
torch.load(create_filename("model", None, args_),
map_location=torch.device("cpu")))
# Train and save
learning_curves = train_model(args, dataset, model, simulator)
# Save
logger.info("Saving model")
torch.save(model.state_dict(), create_filename("model", None, args))
np.save(create_filename("learning_curve", None, args), learning_curves)
logger.info("All done! Have a nice day!")
distances_gen = simulator.distance_from_manifold(x_gen)
mean_gen_distance = np.mean(distances_gen)
# Report results
logger.info("Results:")
logger.info(" test log p: %s", mean_log_likelihood_test)
logger.info(" test reco err: %s", mean_reco_error_test)
logger.info(" gen distance: %s", mean_gen_distance)
return (-1.0 * margs.metricnllfactor * mean_log_likelihood_test +
margs.metricrecoerrorfactor * mean_reco_error_test +
margs.metricdistancefactor * mean_gen_distance)
# Load saved study object
if args.resumestudy:
filename = create_filename("paramscan", None, args)
logger.info("Loading parameter scan from %s", filename)
with open(filename, "rb") as file:
study = pickle.load(file)
# Optimize!
study = optuna.create_study(study_name=args.paramscanstudyname,
direction="minimize")
try:
study.optimize(objective, n_trials=args.trials)
except (KeyboardInterrupt, SystemExit):
logger.warning("Optimization interrupted!")
# Report best results
logger.info("Best parameters:")
create_modelname(args)
logger.info("Evaluating simulator truth")
else:
create_modelname(args)
logger.info("Evaluating model %s", args.modelname)
# Bug fix related to some num_workers > 1 and CUDA. Bad things happen otherwise!
torch.multiprocessing.set_start_method("spawn", force=True)
# Data set
simulator = load_simulator(args)
# Load model
if not args.truth:
model = create_model(args, simulator=simulator)
model.load_state_dict(torch.load(create_filename("model", None, args), map_location=torch.device("cpu")))
model.eval()
else:
model = None
# Evaluate generative performance
if args.skipgeneration:
logger.info("Skipping generative evaluation as per request.")
elif not args.truth:
x_gen = sample_from_model(args, model, simulator)
evaluate_model_samples(args, simulator, x_gen)
if args.skipinference:
logger.info("Skipping all inference tasks as per request. Have a nice day!")
exit()