<filename>scripts/python/combine_models.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2020  Johns Hopkins University (Author: <NAME>)
# Apache 2.0

import argparse
import sys
import os
import json
from itertools import chain

import css.models as models
import css.objectives as objectives

import torch


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "omodel",
        help="path to output model",
        type=str,
    )
    parser.add_argument("conf", type=str)
    parser.add_argument("--save-models", action="store_true")
    parser.add_argument("--models", nargs="+", type=str, help="paths to models")
    args = parser.parse_args()

    conf = json.load(open(args.conf))
    new_model = models.MODELS[conf["model"]].build_model(conf)
    objective = objectives.OBJECTIVES[conf["objective"]].build_objective(conf)

    params = list(
        filter(
            lambda p: p.requires_grad,
            chain(new_model.parameters(), objective.parameters()),
        )
    )

    optimizers = {
        "sgd": torch.optim.SGD(params, lr=conf["lr"], momentum=0.0),
        "adadelta": torch.optim.Adadelta(params, lr=conf["lr"]),
        "adam": torch.optim.Adam(
            params, lr=conf["lr"], weight_decay=conf["weight_decay"]
        ),
    }

    optimizer = optimizers[conf["optim"]]

    new_mdl_dict = new_model.state_dict()
    new_optim_dict = optimizer.state_dict()
    new_objective_dict = objective.state_dict()

    for name, param in new_mdl_dict.items():
        if len(param.size()) > 0:
            param.mul_(0.0)

    # Some requested models may not exist -- for instance we only save models
    # that completed successfully.
    model_list = [m for m in args.models if os.path.exists(m)]
    try:
        fraction = 1.0 / (len(model_list))
    except ZeroDivisionError:
        print("No remaining models", file=sys.stderr)
        sys.exit(1)

    for i, m in enumerate(model_list):
        print("Combining Model ", i, " : ", m, " ...")
        state_dict = torch.load(m, map_location=torch.device("cpu"))
        # ----------------------- Model -------------------------
        # To combine models, we just average the weights
        for name, p in state_dict["model"].items():
            if name in new_mdl_dict:
                if len(p.size()) != 0:
                    new_mdl_dict[name].add_(p, alpha=fraction)
                else:
                    new_mdl_dict[name] = (p * fraction).type(new_mdl_dict[name].dtype)

        # --------------------- Objectives ---------------------
        # To combine objectives is harder: We average parameter weights if
        # applicable, but in the case of some models such as the EBM models
        # we have to specify how to combine things like the sampling buffer.
        # This combination is model specific and should therefore written as a
        # as method in the objective's class. For now we have just done it
        # here though.
        update_opt_state_dict(new_optim_dict, state_dict["optimizer"], fraction)
        new_objective_dict = objective.add_state_dict(
            new_objective_dict,
            state_dict["objective"],
            fraction,
            iteration=i,
        )

    new_state_dict = {
        "model": new_mdl_dict,
        "objective": new_objective_dict,
        "optimizer": state_dict["optimizer"],
        "lr_sched": state_dict["lr_sched"],
        "epoch": state_dict["epoch"],
    }

    torch.save(
        new_state_dict,
        args.omodel,
    )

    if not args.save_models:
        for m in model_list:
            os.remove(m)


def update_opt_state_dict(state_dict1, state_dict2, fraction):
    """
    Update state_dict1, with state_dict2 where values are
    val1 + fraction*val2
    """
    groups2 = state_dict2["param_groups"]
    groups1 = state_dict1["param_groups"]

    if len(groups1) != len(groups2):
        raise ValueError("state dict as a different number of parameter groups")

    param_lens = (len(g["params"]) for g in groups1)
    saved_lens = (len(g["params"]) for g in groups2)
    if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
        raise ValueError(
            "loaded state dict contains a parameter group that "
            "doesn't match the size of the optimizer's group"
        )

    id_map = {
        p: old_id
        for old_id, p in zip(
            chain(*(g["params"] for g in groups1)),
            chain(*(g["params"] for g in groups2)),
        )
    }

    for k, v in state_dict2["state"].items():
        if k in id_map:
            param = id_map[k]
            if param in state_dict1["state"]:
                for p_name, p in v.items():
                    if isinstance(p, torch.Tensor):
                        state_dict1["state"][param][p_name] += fraction * p
            else:
                state_dict1["state"][param] = {
                    key: fraction * val for key, val in v.items()
                }
        else:
            for p_name, p in v.items():
                if isinstance(p, torch.Tensor):
                    state_dict1["state"][k][p_name] = fraction * p


if __name__ == "__main__":
    main()
