vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/gim/gluefactory/train.py

@@ -0,0 +1,691 @@
+"""
+A generic training script that works with any model and dataset.
+
+Author: Paul-Edouard Sarlin (skydes)
+"""
+
+import argparse
+import copy
+import re
+import shutil
+import signal
+from collections import defaultdict
+from pathlib import Path
+from pydoc import locate
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from torch.cuda.amp import GradScaler, autocast
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from . import __module_name__, logger
+from .datasets import get_dataset
+from .eval import run_benchmark
+from .models import get_model
+from .settings import EVAL_PATH, TRAINING_PATH
+from .utils.experiments import get_best_checkpoint, get_last_checkpoint, save_experiment
+from .utils.stdout_capturing import capture_outputs
+from .utils.tensor import batch_to_device
+from .utils.tools import (
+    AverageMetric,
+    MedianMetric,
+    PRMetric,
+    RecallMetric,
+    fork_rng,
+    set_seed,
+)
+
+# @TODO: Fix pbar pollution in logs
+# @TODO: add plotting during evaluation
+
+default_train_conf = {
+    "seed": "???",  # training seed
+    "epochs": 1,  # number of epochs
+    "optimizer": "adam",  # name of optimizer in [adam, sgd, rmsprop]
+    "opt_regexp": None,  # regular expression to filter parameters to optimize
+    "optimizer_options": {},  # optional arguments passed to the optimizer
+    "lr": 0.001,  # learning rate
+    "lr_schedule": {
+        "type": None,  # string in {factor, exp, member of torch.optim.lr_scheduler}
+        "start": 0,
+        "exp_div_10": 0,
+        "on_epoch": False,
+        "factor": 1.0,
+        "options": {},  # add lr_scheduler arguments here
+    },
+    "lr_scaling": [(100, ["dampingnet.const"])],
+    "eval_every_iter": 1000,  # interval for evaluation on the validation set
+    "save_every_iter": 5000,  # interval for saving the current checkpoint
+    "log_every_iter": 200,  # interval for logging the loss to the console
+    "log_grad_every_iter": None,  # interval for logging gradient hists
+    "test_every_epoch": 1,  # interval for evaluation on the test benchmarks
+    "keep_last_checkpoints": 10,  # keep only the last X checkpoints
+    "load_experiment": None,  # initialize the model from a previous experiment
+    "median_metrics": [],  # add the median of some metrics
+    "recall_metrics": {},  # add the recall of some metrics
+    "pr_metrics": {},  # add pr curves, set labels/predictions/mask keys
+    "best_key": "loss/total",  # key to use to select the best checkpoint
+    "dataset_callback_fn": None,  # data func called at the start of each epoch
+    "dataset_callback_on_val": False,  # call data func on val data?
+    "clip_grad": None,
+    "pr_curves": {},
+    "plot": None,
+    "submodules": [],
+}
+default_train_conf = OmegaConf.create(default_train_conf)
+
+
+@torch.no_grad()
+def do_evaluation(model, loader, device, loss_fn, conf, pbar=True):
+    model.eval()
+    results = {}
+    pr_metrics = defaultdict(PRMetric)
+    figures = []
+    if conf.plot is not None:
+        n, plot_fn = conf.plot
+        plot_ids = np.random.choice(len(loader), min(len(loader), n), replace=False)
+    for i, data in enumerate(
+        tqdm(loader, desc="Evaluation", ascii=True, disable=not pbar)
+    ):
+        data = batch_to_device(data, device, non_blocking=True)
+        with torch.no_grad():
+            pred = model(data)
+            losses, metrics = loss_fn(pred, data)
+            if conf.plot is not None and i in plot_ids:
+                figures.append(locate(plot_fn)(pred, data))
+            # add PR curves
+            for k, v in conf.pr_curves.items():
+                pr_metrics[k].update(
+                    pred[v["labels"]],
+                    pred[v["predictions"]],
+                    mask=pred[v["mask"]] if "mask" in v.keys() else None,
+                )
+            del pred, data
+        numbers = {**metrics, **{"loss/" + k: v for k, v in losses.items()}}
+        for k, v in numbers.items():
+            if k not in results:
+                results[k] = AverageMetric()
+                if k in conf.median_metrics:
+                    results[k + "_median"] = MedianMetric()
+                if k in conf.recall_metrics.keys():
+                    q = conf.recall_metrics[k]
+                    results[k + f"_recall{int(q)}"] = RecallMetric(q)
+            results[k].update(v)
+            if k in conf.median_metrics:
+                results[k + "_median"].update(v)
+            if k in conf.recall_metrics.keys():
+                q = conf.recall_metrics[k]
+                results[k + f"_recall{int(q)}"].update(v)
+        del numbers
+    results = {k: results[k].compute() for k in results}
+    return results, {k: v.compute() for k, v in pr_metrics.items()}, figures
+
+
+def filter_parameters(params, regexp):
+    """Filter trainable parameters based on regular expressions."""
+
+    # Examples of regexp:
+    #     '.*(weight|bias)$'
+    #     'cnn\.(enc0|enc1).*bias'
+    def filter_fn(x):
+        n, p = x
+        match = re.search(regexp, n)
+        if not match:
+            p.requires_grad = False
+        return match
+
+    params = list(filter(filter_fn, params))
+    assert len(params) > 0, regexp
+    logger.info("Selected parameters:\n" + "\n".join(n for n, p in params))
+    return params
+
+
+def get_lr_scheduler(optimizer, conf):
+    """Get lr scheduler specified by conf.train.lr_schedule."""
+    if conf.type not in ["factor", "exp", None]:
+        return getattr(torch.optim.lr_scheduler, conf.type)(optimizer, **conf.options)
+
+    # backward compatibility
+    def lr_fn(it):  # noqa: E306
+        if conf.type is None:
+            return 1
+        if conf.type == "factor":
+            return 1.0 if it < conf.start else conf.factor
+        if conf.type == "exp":
+            gam = 10 ** (-1 / conf.exp_div_10)
+            return 1.0 if it < conf.start else gam
+        else:
+            raise ValueError(conf.type)
+
+    return torch.optim.lr_scheduler.MultiplicativeLR(optimizer, lr_fn)
+
+
+def pack_lr_parameters(params, base_lr, lr_scaling):
+    """Pack each group of parameters with the respective scaled learning rate."""
+    filters, scales = tuple(zip(*[(n, s) for s, names in lr_scaling for n in names]))
+    scale2params = defaultdict(list)
+    for n, p in params:
+        scale = 1
+        # TODO: use proper regexp rather than just this inclusion check
+        is_match = [f in n for f in filters]
+        if any(is_match):
+            scale = scales[is_match.index(True)]
+        scale2params[scale].append((n, p))
+    logger.info(
+        "Parameters with scaled learning rate:\n%s",
+        {s: [n for n, _ in ps] for s, ps in scale2params.items() if s != 1},
+    )
+    lr_params = [
+        {"lr": scale * base_lr, "params": [p for _, p in ps]}
+        for scale, ps in scale2params.items()
+    ]
+    return lr_params
+
+
+def training(rank, conf, output_dir, args):
+    if args.restore:
+        logger.info(f"Restoring from previous training of {args.experiment}")
+        try:
+            init_cp = get_last_checkpoint(args.experiment, allow_interrupted=False)
+        except AssertionError:
+            init_cp = get_best_checkpoint(args.experiment)
+        logger.info(f"Restoring from checkpoint {init_cp.name}")
+        init_cp = torch.load(str(init_cp), map_location="cpu")
+        conf = OmegaConf.merge(OmegaConf.create(init_cp["conf"]), conf)
+        conf.train = OmegaConf.merge(default_train_conf, conf.train)
+        epoch = init_cp["epoch"] + 1
+
+        # get the best loss or eval metric from the previous best checkpoint
+        best_cp = get_best_checkpoint(args.experiment)
+        best_cp = torch.load(str(best_cp), map_location="cpu")
+        best_eval = best_cp["eval"][conf.train.best_key]
+        del best_cp
+    else:
+        # we start a new, fresh training
+        conf.train = OmegaConf.merge(default_train_conf, conf.train)
+        epoch = 0
+        best_eval = float("inf")
+        if conf.train.load_experiment:
+            logger.info(f"Will fine-tune from weights of {conf.train.load_experiment}")
+            # the user has to make sure that the weights are compatible
+            try:
+                init_cp = get_last_checkpoint(conf.train.load_experiment)
+            except AssertionError:
+                init_cp = get_best_checkpoint(conf.train.load_experiment)
+            # init_cp = get_last_checkpoint(conf.train.load_experiment)
+            init_cp = torch.load(str(init_cp), map_location="cpu")
+            # load the model config of the old setup, and overwrite with current config
+            conf.model = OmegaConf.merge(
+                OmegaConf.create(init_cp["conf"]).model, conf.model
+            )
+            print(conf.model)
+        else:
+            init_cp = None
+
+    OmegaConf.set_struct(conf, True)  # prevent access to unknown entries
+    set_seed(conf.train.seed)
+    if rank == 0:
+        writer = SummaryWriter(log_dir=str(output_dir))
+
+    data_conf = copy.deepcopy(conf.data)
+    if args.distributed:
+        logger.info(f"Training in distributed mode with {args.n_gpus} GPUs")
+        assert torch.cuda.is_available()
+        device = rank
+        torch.distributed.init_process_group(
+            backend="nccl",
+            world_size=args.n_gpus,
+            rank=device,
+            init_method="file://" + str(args.lock_file),
+        )
+        torch.cuda.set_device(device)
+
+        # adjust batch size and num of workers since these are per GPU
+        if "batch_size" in data_conf:
+            data_conf.batch_size = int(data_conf.batch_size / args.n_gpus)
+        if "train_batch_size" in data_conf:
+            data_conf.train_batch_size = int(data_conf.train_batch_size / args.n_gpus)
+        if "num_workers" in data_conf:
+            data_conf.num_workers = int(
+                (data_conf.num_workers + args.n_gpus - 1) / args.n_gpus
+            )
+    else:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device {device}")
+
+    dataset = get_dataset(data_conf.name)(data_conf)
+
+    # Optionally load a different validation dataset than the training one
+    val_data_conf = conf.get("data_val", None)
+    if val_data_conf is None:
+        val_dataset = dataset
+    else:
+        val_dataset = get_dataset(val_data_conf.name)(val_data_conf)
+
+    # @TODO: add test data loader
+
+    if args.overfit:
+        # we train and eval with the same single training batch
+        logger.info("Data in overfitting mode")
+        assert not args.distributed
+        train_loader = dataset.get_overfit_loader("train")
+        val_loader = val_dataset.get_overfit_loader("val")
+    else:
+        train_loader = dataset.get_data_loader("train", distributed=args.distributed)
+        val_loader = val_dataset.get_data_loader("val")
+    if rank == 0:
+        logger.info(f"Training loader has {len(train_loader)} batches")
+        logger.info(f"Validation loader has {len(val_loader)} batches")
+
+    # interrupts are caught and delayed for graceful termination
+    def sigint_handler(signal, frame):
+        logger.info("Caught keyboard interrupt signal, will terminate")
+        nonlocal stop
+        if stop:
+            raise KeyboardInterrupt
+        stop = True
+
+    stop = False
+    signal.signal(signal.SIGINT, sigint_handler)
+    model = get_model(conf.model.name)(conf.model).to(device)
+    if args.compile:
+        model = torch.compile(model, mode=args.compile)
+    loss_fn = model.loss
+    if init_cp is not None:
+        model.load_state_dict(init_cp["model"], strict=False)
+    if args.distributed:
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device])
+    if rank == 0 and args.print_arch:
+        logger.info(f"Model: \n{model}")
+
+    torch.backends.cudnn.benchmark = True
+    if args.detect_anomaly:
+        torch.autograd.set_detect_anomaly(True)
+
+    optimizer_fn = {
+        "sgd": torch.optim.SGD,
+        "adam": torch.optim.Adam,
+        "adamw": torch.optim.AdamW,
+        "rmsprop": torch.optim.RMSprop,
+    }[conf.train.optimizer]
+    params = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
+    if conf.train.opt_regexp:
+        params = filter_parameters(params, conf.train.opt_regexp)
+    all_params = [p for n, p in params]
+
+    lr_params = pack_lr_parameters(params, conf.train.lr, conf.train.lr_scaling)
+    optimizer = optimizer_fn(
+        lr_params, lr=conf.train.lr, **conf.train.optimizer_options
+    )
+    scaler = GradScaler(enabled=args.mixed_precision is not None)
+    logger.info(f"Training with mixed_precision={args.mixed_precision}")
+
+    mp_dtype = {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        None: torch.float32,  # we disable it anyway
+    }[args.mixed_precision]
+
+    results = None  # fix bug with it saving
+
+    lr_scheduler = get_lr_scheduler(optimizer=optimizer, conf=conf.train.lr_schedule)
+    if args.restore:
+        optimizer.load_state_dict(init_cp["optimizer"])
+        if "lr_scheduler" in init_cp:
+            lr_scheduler.load_state_dict(init_cp["lr_scheduler"])
+
+    if rank == 0:
+        logger.info(
+            "Starting training with configuration:\n%s", OmegaConf.to_yaml(conf)
+        )
+    losses_ = None
+
+    def trace_handler(p):
+        # torch.profiler.tensorboard_trace_handler(str(output_dir))
+        output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
+        print(output)
+        p.export_chrome_trace("trace_" + str(p.step_num) + ".json")
+        p.export_stacks("/tmp/profiler_stacks.txt", "self_cuda_time_total")
+
+    if args.profile:
+        prof = torch.profiler.profile(
+            schedule=torch.profiler.schedule(wait=1, warmup=1, active=1, repeat=1),
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(str(output_dir)),
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=True,
+        )
+        prof.__enter__()
+    while epoch < conf.train.epochs and not stop:
+        if rank == 0:
+            logger.info(f"Starting epoch {epoch}")
+
+        # we first run the eval
+        if (
+            rank == 0
+            and epoch % conf.train.test_every_epoch == 0
+            and args.run_benchmarks
+        ):
+            for bname, eval_conf in conf.get("benchmarks", {}).items():
+                logger.info(f"Running eval on {bname}")
+                s, f, r = run_benchmark(
+                    bname,
+                    eval_conf,
+                    EVAL_PATH / bname / args.experiment / str(epoch),
+                    model.eval(),
+                )
+                logger.info(str(s))
+                for metric_name, value in s.items():
+                    writer.add_scalar(f"test/{bname}/{metric_name}", value, epoch)
+                for fig_name, fig in f.items():
+                    writer.add_figure(f"figures/{bname}/{fig_name}", fig, epoch)
+
+        # set the seed
+        set_seed(conf.train.seed + epoch)
+
+        # update learning rate
+        if conf.train.lr_schedule.on_epoch and epoch > 0:
+            old_lr = optimizer.param_groups[0]["lr"]
+            lr_scheduler.step()
+            logger.info(
+                f'lr changed from {old_lr} to {optimizer.param_groups[0]["lr"]}'
+            )
+        if args.distributed:
+            train_loader.sampler.set_epoch(epoch)
+        if epoch > 0 and conf.train.dataset_callback_fn and not args.overfit:
+            loaders = [train_loader]
+            if conf.train.dataset_callback_on_val:
+                loaders += [val_loader]
+            for loader in loaders:
+                if isinstance(loader.dataset, torch.utils.data.Subset):
+                    getattr(loader.dataset.dataset, conf.train.dataset_callback_fn)(
+                        conf.train.seed + epoch
+                    )
+                else:
+                    getattr(loader.dataset, conf.train.dataset_callback_fn)(
+                        conf.train.seed + epoch
+                    )
+        for it, data in enumerate(train_loader):
+            tot_it = (len(train_loader) * epoch + it) * (
+                args.n_gpus if args.distributed else 1
+            )
+            tot_n_samples = tot_it
+            if not args.log_it:
+                # We normalize the x-axis of tensorflow to num samples!
+                tot_n_samples *= train_loader.batch_size
+
+            model.train()
+            optimizer.zero_grad()
+
+            with autocast(enabled=args.mixed_precision is not None, dtype=mp_dtype):
+                data = batch_to_device(data, device, non_blocking=True)
+                pred = model(data)
+                losses, _ = loss_fn(pred, data)
+                loss = torch.mean(losses["total"])
+            if torch.isnan(loss).any():
+                print(f"Detected NAN, skipping iteration {it}")
+                del pred, data, loss, losses
+                continue
+
+            do_backward = loss.requires_grad
+            if args.distributed:
+                do_backward = torch.tensor(do_backward).float().to(device)
+                torch.distributed.all_reduce(
+                    do_backward, torch.distributed.ReduceOp.PRODUCT
+                )
+                do_backward = do_backward > 0
+            if do_backward:
+                scaler.scale(loss).backward()
+                if args.detect_anomaly:
+                    # Check for params without any gradient which causes
+                    # problems in distributed training with checkpointing
+                    detected_anomaly = False
+                    for name, param in model.named_parameters():
+                        if param.grad is None and param.requires_grad:
+                            print(f"param {name} has no gradient.")
+                            detected_anomaly = True
+                    if detected_anomaly:
+                        raise RuntimeError("Detected anomaly in training.")
+                if conf.train.get("clip_grad", None):
+                    scaler.unscale_(optimizer)
+                    try:
+                        torch.nn.utils.clip_grad_norm_(
+                            all_params,
+                            max_norm=conf.train.clip_grad,
+                            error_if_nonfinite=True,
+                        )
+                        scaler.step(optimizer)
+                    except RuntimeError:
+                        logger.warning("NaN detected in gradients. Skipping iteration.")
+                    scaler.update()
+                else:
+                    scaler.step(optimizer)
+                    scaler.update()
+                if not conf.train.lr_schedule.on_epoch:
+                    lr_scheduler.step()
+            else:
+                if rank == 0:
+                    logger.warning(f"Skip iteration {it} due to detach.")
+
+            if args.profile:
+                prof.step()
+
+            if it % conf.train.log_every_iter == 0:
+                for k in sorted(losses.keys()):
+                    if args.distributed:
+                        losses[k] = losses[k].sum(-1)
+                        torch.distributed.reduce(losses[k], dst=0)
+                        losses[k] /= train_loader.batch_size * args.n_gpus
+                    losses[k] = torch.mean(losses[k], -1)
+                    losses[k] = losses[k].item()
+                if rank == 0:
+                    str_losses = [f"{k} {v:.3E}" for k, v in losses.items()]
+                    logger.info(
+                        "[E {} | it {}] loss {{{}}}".format(
+                            epoch, it, ", ".join(str_losses)
+                        )
+                    )
+                    for k, v in losses.items():
+                        writer.add_scalar("training/" + k, v, tot_n_samples)
+                    writer.add_scalar(
+                        "training/lr", optimizer.param_groups[0]["lr"], tot_n_samples
+                    )
+                    writer.add_scalar("training/epoch", epoch, tot_n_samples)
+
+            if conf.train.log_grad_every_iter is not None:
+                if it % conf.train.log_grad_every_iter == 0:
+                    grad_txt = ""
+                    for name, param in model.named_parameters():
+                        if param.grad is not None and param.requires_grad:
+                            if name.endswith("bias"):
+                                continue
+                            writer.add_histogram(
+                                f"grad/{name}", param.grad.detach(), tot_n_samples
+                            )
+                            norm = torch.norm(param.grad.detach(), 2)
+                            grad_txt += f"{name} {norm.item():.3f}  \n"
+                    writer.add_text("grad/summary", grad_txt, tot_n_samples)
+            del pred, data, loss, losses
+
+            # Run validation
+            if (
+                (
+                    it % conf.train.eval_every_iter == 0
+                    and (it > 0 or epoch == -int(args.no_eval_0))
+                )
+                or stop
+                or it == (len(train_loader) - 1)
+            ):
+                with fork_rng(seed=conf.train.seed):
+                    results, pr_metrics, figures = do_evaluation(
+                        model,
+                        val_loader,
+                        device,
+                        loss_fn,
+                        conf.train,
+                        pbar=(rank == -1),
+                    )
+
+                if rank == 0:
+                    str_results = [
+                        f"{k} {v:.3E}"
+                        for k, v in results.items()
+                        if isinstance(v, float)
+                    ]
+                    logger.info(f'[Validation] {{{", ".join(str_results)}}}')
+                    for k, v in results.items():
+                        if isinstance(v, dict):
+                            writer.add_scalars(f"figure/val/{k}", v, tot_n_samples)
+                        else:
+                            writer.add_scalar("val/" + k, v, tot_n_samples)
+                    for k, v in pr_metrics.items():
+                        writer.add_pr_curve("val/" + k, *v, tot_n_samples)
+                    # @TODO: optional always save checkpoint
+                    if results[conf.train.best_key] < best_eval:
+                        best_eval = results[conf.train.best_key]
+                        save_experiment(
+                            model,
+                            optimizer,
+                            lr_scheduler,
+                            conf,
+                            losses_,
+                            results,
+                            best_eval,
+                            epoch,
+                            tot_it,
+                            output_dir,
+                            stop,
+                            args.distributed,
+                            cp_name="checkpoint_best.tar",
+                        )
+                        logger.info(f"New best val: {conf.train.best_key}={best_eval}")
+                    if len(figures) > 0:
+                        for i, figs in enumerate(figures):
+                            for name, fig in figs.items():
+                                writer.add_figure(
+                                    f"figures/{i}_{name}", fig, tot_n_samples
+                                )
+                torch.cuda.empty_cache()  # should be cleared at the first iter
+
+            if (tot_it % conf.train.save_every_iter == 0 and tot_it > 0) and rank == 0:
+                if results is None:
+                    results, _, _ = do_evaluation(
+                        model,
+                        val_loader,
+                        device,
+                        loss_fn,
+                        conf.train,
+                        pbar=(rank == -1),
+                    )
+                    best_eval = results[conf.train.best_key]
+                best_eval = save_experiment(
+                    model,
+                    optimizer,
+                    lr_scheduler,
+                    conf,
+                    losses_,
+                    results,
+                    best_eval,
+                    epoch,
+                    tot_it,
+                    output_dir,
+                    stop,
+                    args.distributed,
+                )
+
+            if stop:
+                break
+
+        if rank == 0:
+            best_eval = save_experiment(
+                model,
+                optimizer,
+                lr_scheduler,
+                conf,
+                losses_,
+                results,
+                best_eval,
+                epoch,
+                tot_it,
+                output_dir=output_dir,
+                stop=stop,
+                distributed=args.distributed,
+            )
+
+        epoch += 1
+
+    logger.info(f"Finished training on process {rank}.")
+    if rank == 0:
+        writer.close()
+
+
+def main_worker(rank, conf, output_dir, args):
+    if rank == 0:
+        with capture_outputs(output_dir / "log.txt"):
+            training(rank, conf, output_dir, args)
+    else:
+        training(rank, conf, output_dir, args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("experiment", type=str)
+    parser.add_argument("--conf", type=str)
+    parser.add_argument(
+        "--mixed_precision",
+        "--mp",
+        default=None,
+        type=str,
+        choices=["float16", "bfloat16"],
+    )
+    parser.add_argument(
+        "--compile",
+        default=None,
+        type=str,
+        choices=["default", "reduce-overhead", "max-autotune"],
+    )
+    parser.add_argument("--overfit", action="store_true")
+    parser.add_argument("--restore", action="store_true")
+    parser.add_argument("--distributed", action="store_true")
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--print_arch", "--pa", action="store_true")
+    parser.add_argument("--detect_anomaly", "--da", action="store_true")
+    parser.add_argument("--log_it", "--log_it", action="store_true")
+    parser.add_argument("--no_eval_0", action="store_true")
+    parser.add_argument("--run_benchmarks", action="store_true")
+    parser.add_argument("dotlist", nargs="*")
+    args = parser.parse_intermixed_args()
+
+    logger.info(f"Starting experiment {args.experiment}")
+    output_dir = Path(TRAINING_PATH, args.experiment)
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    conf = OmegaConf.from_cli(args.dotlist)
+    if args.conf:
+        conf = OmegaConf.merge(OmegaConf.load(args.conf), conf)
+    elif args.restore:
+        restore_conf = OmegaConf.load(output_dir / "config.yaml")
+        conf = OmegaConf.merge(restore_conf, conf)
+    if not args.restore:
+        if conf.train.seed is None:
+            conf.train.seed = torch.initial_seed() & (2**32 - 1)
+        OmegaConf.save(conf, str(output_dir / "config.yaml"))
+
+    # copy gluefactory and submodule into output dir
+    for module in conf.train.get("submodules", []) + [__module_name__]:
+        mod_dir = Path(__import__(str(module)).__file__).parent
+        shutil.copytree(mod_dir, output_dir / module, dirs_exist_ok=True)
+
+    if args.distributed:
+        args.n_gpus = torch.cuda.device_count()
+        args.lock_file = output_dir / "distributed_lock"
+        if args.lock_file.exists():
+            args.lock_file.unlink()
+        torch.multiprocessing.spawn(
+            main_worker, nprocs=args.n_gpus, args=(conf, output_dir, args)
+        )
+    else:
+        main_worker(0, conf, output_dir, args)