# -- Path Setup ----------------------------------------------------------------------

# if extensions (or modules to document with autodoc) are in another directory add these directories
# to sys.path here. If the directory is relative to the documentation root, use
# os.path.abspath to make it absolute
import os
import sys
# sys.path.insert(0, os.path.abspath('../../'))
# print("CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======CWD======S")
# print(os.getcwd())
print("os.path.abspath(os.path.join(__file__, ../..))",os.path.abspath( os.path.join(__file__, "../../..")
))
sys.path.insert(0, os.path.abspath('..'))
sys.path.append(os.path.abspath(
    os.path.join(__file__, "../../..")
))


# Configuration file for the Sphinx documentation builder.

# -- Project information

project = 'push'
copyright = 'The Push authors'
author = 'Daniel Huang'

release = '0.1'
version = '0.1.0'

# -- General configuration

extensions = [
    'sphinx.ext.duration',
    'sphinx.ext.doctest',
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.intersphinx',
    "nbsphinx"
]
autodoc_mock_imports = ["tqdm","numpy"]

intersphinx_mapping = {
    'python': ('https://docs.python.org/3/', None),
    'sphinx': ('https://www.sphinx-doc.org/en/master/', None),
}
intersphinx_disabled_domains = ['std']

templates_path = ['_templates']

# -- Options for HTML output

html_theme = 'sphinx_rtd_theme'

# -- Options for EPUB output
epub_show_urls = 'footnote'
autosummary_generate = True  # Turn on sphinx.ext.autosummary
docs/source/conf.py
lbai-lab/push
import argparse
from typing import *
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import push.bayes.ensemble
import push.bayes.stein_vgd
import push.bayes.swag

import time


torch.manual_seed(34)


"""[Note]

Quick Usage

1. `python test_basic.py -m ensemble`
2. `python test_basic.py -m stein_vgd`
3. `python test_basic.py -m mswag`

"""

# =============================================================================
# Simple Dataset + Neural Network
# =============================================================================

class RandDataset(Dataset):
    def __init__(self, batch_size, N, D):
        self.xs = torch.randn(batch_size*N, D)
        self.ys = torch.randn(batch_size*N, 1)

    def __len__(self):
        return len(self.xs)

    def __getitem__(self, idx):
        return self.xs[idx], self.ys[idx]


class MiniNN(nn.Module):
    def __init__(self, D):
        super(MiniNN, self).__init__()
        self.fc1 = nn.Linear(D, D)
        self.fc2 = nn.Linear(D, D)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.ReLU()(x)
        x = self.fc2(x)
        return x
    

class BiggerNN(nn.Module):
    def __init__(self, n, D):
        super(BiggerNN, self).__init__()
        self.minis = []
        self.n = n
        for i in range(0, n):
            self.minis += [MiniNN(D)]
            self.add_module("mini_layer"+str(i), self.minis[-1])
        self.fc = nn.Linear(D, 1)
            
    def forward(self, x):
        for i in range(0, self.n):
            x = self.minis[i](x)
        return self.fc(x)


# =============================================================================
# Basic Testing
# =============================================================================

if __name__ == "__main__":
    # Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-L", "--layers", type=int, default=2,
                        help="size of neural network")
    parser.add_argument("-D", "--dimension", type=int, default=20,
                        help="size of neural network")
    parser.add_argument("-p", "--particles", type=int, default=2,
                        help="number of particles for SVGD or mswag")
    parser.add_argument("-b", "--batchsize", type=int, default=128,
                        help="number of batches")
    parser.add_argument("-e", "--epochs", type=int, default=10,
                        help="number of batches")
    parser.add_argument("-n", "--size", type=int, default=10,
                        help="size of dataset")
    parser.add_argument("-d", "--devices", type=int, default=1,
                        help="number of devices")
    parser.add_argument("-cs", "--cache_size", type=int, default=4,
                        help="size of the cache")
    parser.add_argument("-vs", "--view_size", type=int, default=4,
                        help="size of the cache")
    parser.add_argument("-m", "--method", type=str, default="ensemble",
                        choices=[
                            "ensemble",
                            "swag-central",
                            "swag",
                            "mswag",
                            "stein_vgd",
                        ],
                        help="bayesian inference method")
    args = parser.parse_args()

    # Dataset
    L = args.layers
    D = args.dimension
    N = args.size
    dataset = RandDataset(args.batchsize, N, D)
    dataloader = DataLoader(dataset, batch_size=args.batchsize, shuffle=True)
    
    # Dispatch
    if args.method == "ensemble":
        def mk_dataloader():
            return DataLoader(RandDataset(args.batchsize, args.size, args.dimension), batch_size=128, shuffle=True)

        epochs = 10
        num_ensembles = 3
        push.bayes.ensemble.train_deep_ensemble(
            dataloader, # mk_dataloader,
            torch.nn.MSELoss(),
            epochs,
            BiggerNN, L, D,
            num_devices=args.devices,
            num_ensembles=num_ensembles
        )

    elif args.method == "mswag":
        pretrain_epochs = 10
        swag_epochs = 5
        cache_size=4
        view_size=4
        start = time.perf_counter()
        push.bayes.swag.train_mswag(
            dataloader,
            torch.nn.MSELoss(),
            pretrain_epochs,
            swag_epochs,
            args.particles,
            cache_size,
            view_size, 
            BiggerNN, L, D
        )
        print("Time elapsed", time.perf_counter() - start)

    elif args.method == "stein_vgd":
        import numpy as np
        start = time.perf_counter()
        print("Size of model", sum([np.prod(p.size()) for p in BiggerNN(L, D).parameters()]))
        print("Size of dataset", len(dataloader.dataset))
        svgd_state = {
            "L": L,
            "D": D,
            "N": N,
        }
        push.bayes.stein_vgd.train_svgd(
            dataloader, torch.nn.MSELoss(),
            args.epochs, args.particles,
            BiggerNN, L, D,
            lengthscale=1.0, lr=1e-3, prior=None,
            cache_size=args.cache_size, view_size=args.view_size, num_devices=args.devices,
            svgd_entry=push.bayes.stein_vgd._svgd_leader, svgd_state=svgd_state
        )
        print("Time elapsed", time.perf_counter() - start)

    else:
        raise ValueError(f"Method {args.method} not supported ...")

docs/source/tutorials/01_Bayesian_Deep_Learning/test_basic.py
lbai-lab/push
import os
import glob
import tqdm
import torch
import pickle
import requests
import argparse
import numpy as np

from torch.utils.data import Dataset


# Parent class of MD17 
class MD17Dataset(Dataset):
    def __init__(self, style, task, split, root="./"):
        # Names
        self.Mdict = {"a":"aspirin", "b":"benzene", "e":"ethanol", "m":"malonaldehyde", "n":"naphthalene", "s":"salicylic", "t":"toluene", "u":"uracil"}
        self.Ndict = {"a":21, "b":12, "e":9, "m":9, "n":18, "s":16, "t":15, "u":12}

        # Download missing data
        missing = [v for v in self.Mdict.values() if not os.path.isfile(f"{root}/{v}_dft.npz")]
        if missing: self.download(root, missing) 
        
        datas_dict, labels_dict = self.load_dataset(root, "*", style, task, split)
        self.datas, self.labels = self.unroll_data(datas_dict, labels_dict)
        self.covariance = self.cov_matrix(self.labels)

    def __len__(self):
        return len(self.datas["R"])

    def __getitem__(self, idx):
        return self.datas["R"][idx], self.datas["z"], self.labels["E"][idx], self.labels["F"][idx], self.labels["covariance_inv"], self.labels["mean"]

    def collate(self, batch):
        data = {"R":[], "z":[], "batch":[], "n":[]}
        label = {"E":[], "F":[], "covariance_inv":batch[0][4], "mean":batch[0][5]}      

        for i, b in enumerate(batch):
            data["R"].append(b[0])
            data["z"].append(b[1])
            data["n"].append(len(b[1]))
            data["batch"].append(torch.ones((b[1].size()), dtype=torch.int64)*i)
            label["E"].append(b[2])
            label["F"].append(b[3])
        data["R"] = torch.cat(data["R"])
        data["z"] = torch.cat(data["z"])
        data["n"] = torch.Tensor(data["n"])
        data["batch"] = torch.cat(data["batch"])
        label["E"] = torch.stack(label["E"])
        label["F"] = torch.stack(label["F"])

        return data, label

    def print_info(self):
        print("Data sizes:")
        for m in self.Mdict.values():
            R = "R"
            print(f"{m}: {len(self.datas[m][R])}")
        return

    def download(self, path, molecules):
        print("Downloading missing datas")
        for m in tqdm.tqdm(molecules):
            if m == "benzene":
                r = requests.get(f"http://www.quantum-machine.org/gdml/data/npz/md17_{m}2017.npz")
            else:
                r = requests.get(f"http://www.quantum-machine.org/gdml/data/npz/md17_{m}.npz")
            open(os.path.join(path, f"{m}_dft.npz") , "wb").write(r.content)

        return 

    def load_dataset(self, root, style, molecule, task, split):
        npzs = glob.glob(f"{root}/{molecule}.npz")
        datas, labels = dict(), dict()

        for npz in npzs:
            z = np.load(npz, allow_pickle=True)
            name = npz.split("/")[-1].split(".")[0]
            datas[name], labels[name] = dict(), dict()

            with open(f"{root}/{split}/{molecule}/{task}.pkl", "rb") as f:
                idx = pickle.load(f)

            # Keys are ['E', 'name', 'F', 'theory', 'R', 'z', 'type', 'md5']
            datas[name]["R"] = z["R"][idx]
            datas[name]["z"] = z["z"]
            labels[name]["E"] = z["E"][idx]
            labels[name]["F"] = z["F"][idx]
        return datas, labels

    def cov_mean_matrix(self, label):
        E = label["E"]
        F = label["F"]
        F = F.reshape((F.size()[0],-1))
        EF = torch.t(torch.cat((E,F), axis=1))

        self.covariance_inv = torch.inverse(torch.cov(EF))
        self.mean = torch.mean(EF, dim=1)
        self.labels["covariance_inv"] = self.covariance_inv
        self.labels["mean"] = self.mean

        return 
    
    def unroll_data(self, datas_dict, labels_dict):
        datas, labels = dict(), dict()
        #TBD

        return datas, labels
    

class MD17SingleDataset(MD17Dataset):
    def __init__(self, style,  m, task, split, root="./"):
        # Initialize
        self.Mdict = {"a":"aspirin", "b":"benzene", "e":"ethanol", "m":"malonaldehyde", "n":"naphthalene", "s":"salicylic", "t":"toluene", "u":"uracil"} # full name of molecule
        self.Ndict = {"a":21, "b":12, "e":9, "m":9, "n":18, "s":16, "t":15, "u":12} # number of atom in each molecule
        self.identifier = f"MD17SingleDataset_{style}_{m}{split}{task}" # name tag for this dataset configuration
        self.m = self.Mdict[m]
        self.natom = self.Ndict[m]
                
        # Load and organized data
        datas_dict, labels_dict = self.load_dataset(root, style, self.m, task, split)
        self.datas, self.labels = self.unroll_data(datas_dict, labels_dict)
        self.cov_mean_matrix(self.labels)

    def print_info(self):
        print("Data sizes:")
        R = "R"
        print(f"{self.m}: {len(self.datas[R])}")
        return

    def unroll_data(self, datas_dict, labels_dict):
        datas = {"R":torch.Tensor(datas_dict[self.m]["R"]), "z":torch.LongTensor(datas_dict[self.m]["z"])} 
        labels = {"E":torch.Tensor(labels_dict[self.m]["E"]), "F":torch.Tensor(labels_dict[self.m]["F"])}

        return datas, labels


def parse_index(exps, l):
    indice = []
    
    for exp in exps.split(","):
        if exp == ":": indice += [i for i in range(l)]
        elif ":" not in exp:  indice += [int(exp)]
        else:
            if exp[0] == ":": indice += [i for i in range(int(exp[1:])+1)]
            elif exp[-1] == ":": indice += [i for i in range(int(exp[:-1])+1)]
            else: indice += [i for i in range(int(exp.split(":")[0]), int(exp.split(":")[1])+1)]

    return indice


def main(args):
    # As a script it could be used to display values in the dataset
    
    # Load dataset
    dataset = MD17SingleDataset(args.style, args.molecule, args.task, args.split)
    if args.info: dataset.print_info()
    
    # Get numerical index
    try:
        indice = parse_index(args.index, len(dataset.labels["E"]))
    except BaseException as err:
        print(err)
        print("expression of index incorrect")

    # Print
    R, z, E, F = "R", "z", "E", "F" 
    #print(f"covariance_inv:{dataset.covariance_inv}")
    #print(len(dataset.datas[z]))
    for idx in indice:
        print(f"{idx:06d}:\n (R={dataset.datas[R][idx]},\t z={dataset.datas[z]}),\t E={dataset.labels[E][idx]},\t F={dataset.labels[F][idx]}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("-r", "--root", type=str, help="path to the root dir of dataset", default="./")
    parser.add_argument("-y", "--style", type=str, default="trajectory", help="specify how to split the data")
    parser.add_argument("-m", "--molecule", type=str, help="lowercase initial of the molecule in the dataset", required=True)
    parser.add_argument("-t", "--task", type=str, help="train or valid or test", required=True)
    parser.add_argument("-s", "--split", type=int, help="the name of dataset subset, aka the number of train samples", required=True)
    parser.add_argument("-i", "--index", type=str, help="indice of the data to be shown", default=":")
    parser.add_argument("-f", "--info", type=bool, help="print the information about this dataset", default=False)

    main(parser.parse_args())

experiments/data/MD17/MD17Dataset.py
lbai-lab/push
from genericpath import isdir
import os
import random
import requests
import pickle
import shutil
import argparse
# from tqdm import tqdm


MOLECULES = {
    "a": "aspirin",
    "b": "benzene",
    "e": "ethanol",
    "m": "malonaldehyde",
    "n": "naphthalene",
    "s": "salicylic",
    "t": "toluene",
    "u": "uracil"
}


def download(path: str) -> None:
    print("Downloading missing data ...")
    missing = [v for v in MOLECULES.values() if not os.path.isfile(f"{path}/{v}.npz")]
    for m in (missing):
        print("m: ", m)
        if m == "benzene":
            r = requests.get(f"http://www.quantum-machine.org/gdml/data/npz/md17_{m}2017.npz")
        else:
            r = requests.get(f"http://www.quantum-machine.org/gdml/data/npz/md17_{m}.npz")
            # print(f"http://www.quantum-machine.org/gdml/data/npz/{m}.npz")
            print(r)
        open(os.path.join(path, f"{m}.npz") , "wb").write(r.content)


def split(args, path):
    sample_num = {
        "aspirin": 211762,
        "benzene": 627983,
        "ethanol": 555092,
        "malonaldehyde": 993237,
        "naphthalene": 326250,
        "salicylic": 320231,
        "toluene": 442790,
        "uracil": 133770
    }

    if os.path.isdir(f"{args.train}"): 
        print("over writing split")
        shutil.rmtree(f"./{args.train}")
    os.mkdir(f"{args.train}")

    for m in sample_num.keys():
        os.mkdir(f"{args.train}/{m}")
        assert args.train+args.valid < sample_num[m]

        pool = [i for i in range(sample_num[m])]
        train_set = random.sample(pool, args.train)
        pool = list(set(pool)-set(train_set))
        valid_set = random.sample(pool, args.valid)
        test_set = list(set(pool)-set(valid_set))
        assert len(train_set)+len(valid_set)+len(test_set) == sample_num[m]

        print(f"{args.train}/{m}/train.pkl")
        with open(f"{args.train}/{m}/train.pkl", "wb") as f:
            pickle.dump(train_set, f)
        with open(f"{args.train}/{m}/valid.pkl", "wb") as f:
            pickle.dump(valid_set, f)
        with open(f"{args.train}/{m}/test.pkl", "wb") as f:
            pickle.dump(test_set, f)

    print("split generated")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--train", type=int, default=1000, help="number of sample in training set")
    parser.add_argument("--valid", type=int, default=100, help="number of sample in validation set")

    args = parser.parse_args()
    download(os.getcwd())
    split(args, "./")

experiments/data/MD17/create_dataset.py
lbai-lab/push
import os
import argparse

from tqdm import tqdm
import pandas as pd
from torchvision.datasets.utils import download_url


def parse_metadata(pde_names):
    """
    This function parses the argument to filter the metadata of files that need to be downloaded.

    Args:
    pde_names: List containing the name of the PDE to be downloaded
    df      : The provided dataframe loaded from the csv file

    Options for pde_names:
    - Advection

    Returns:
    pde_df : Filtered dataframe containing metadata of files to be downloaded
    """

    meta_df = pd.read_csv("download_metadata.csv")

    # Ensure the pde_name is defined
    pde_list = [
        "advection",
        
    ]

    assert all(
        [name.lower() in pde_list for name in pde_names]
    ), "PDE name not defined."

    # Filter the files to be downloaded
    meta_df["PDE"] = meta_df["PDE"].str.lower()
    pde_df = meta_df[meta_df["PDE"].isin(pde_names)]

    return pde_df


def download_data(root_folder, pde_name):
    """ "
    Download data splits specific to a given PDE.

    Args:
    root_folder: The root folder where the data will be downloaded
    pde_name   : The name of the PDE for which the data to be downloaded
    """

    print(f"Downloading data for {pde_name} ...")

    # Load and parse metadata csv file
    pde_df = parse_metadata(pde_name)

    # Iterate filtered dataframe and download the files
    for index, row in tqdm(pde_df.iterrows(), total=pde_df.shape[0]):
        file_path = os.path.join(root_folder, row["Path"])
        download_url(row["URL"], file_path, row["Filename"], md5=row["MD5"])


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(
        prog="Download Script",
        description="Helper script to download the PDEBench datasets",
        epilog="",
    )

    arg_parser.add_argument(
        "--root_folder",
        type=str,
        required=True,
        help="Root folder where the data will be downloaded",
    )
    arg_parser.add_argument(
        "--pde_name",
        action="append",
        help="Name of the PDE dataset to download. You can use this flag multiple times to download multiple datasets",
    )

    args = arg_parser.parse_args()

    download_data(args.root_folder, args.pde_name)
experiments/data/download_direct.py
lbai-lab/push
import os
from torchvision.transforms import ToTensor
import torchvision.transforms as transforms
import torchvision.datasets as datasets


if __name__ == "__main__":
    train_data = datasets.MNIST(
        root = os.getcwd(),
        train = True,                         
        transform = ToTensor(), 
        download = True,            
    )
    test_data = datasets.MNIST(
        root = os.getcwd(), 
        train = False, 
        transform = ToTensor(),
    )
  
# Dataset


experiments/data/vision/create_dataset.py
lbai-lab/push
import argparse

"""[Note]

python gen_train_script.py --model all -d 4 -cp
"""


def generate_one(args, model, f, devices):
    methods = ["ensemble_push", "mswag_push", "svgd_push"]
    if devices == 4:
        particles = [4, 8, 16, 32]
    elif devices == 2:
        particles = [2, 4, 8, 16]
    elif devices == 1:
        particles = [1, 2, 4, 8]

    cp = ""
    if model == "schnet" or model == "cgcnn":
        dl = 800
    elif model == "resnet" or model == "transformer":
        dl = 5120
    elif model == "unet" or model == "fno":
        dl = 2000
        if args.cloud_path:
            cp = "-cp"
        
    for method in methods:
        for p in particles:
            if method == "ensemble_push":
                f.write(f"python train.py -wb -g {args.group} --model {model} -t {method} -n {p} -d {devices} -e {args.epochs} -dl {dl} {cp} || true\n")
            elif method == "mswag_push":
                f.write(f"python train.py -wb -g {args.group} --model {model} -t {method} -n {p} -d {devices} --pretrain_epochs 3 --swag_epochs {args.epochs} -dl {dl} {cp} || true\n")
            else:
                f.write(f"python train.py -wb -g {args.group} --model {model} -t {method} -n {p} -d {devices} -e {args.epochs} -dl {dl} {cp} -mef || true\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-model", "--model", type=str, default="vit",
                        choices=[
                            "all",
                            "schnet",
                            "cgcnn",
                            "unet",
                            "fno",
                            "resnet",
                            "cnn",
                            "transformer"
                        ])
    parser.add_argument("-g", "--group", type=str, default='default')
    parser.add_argument("-e", "--epochs", type=int, default=10)
    parser.add_argument("-d", "--devices", type=int, default=4, choices=[1, 2, 4])
    parser.add_argument("-cp", "--cloud_path", action='store_true')
    args = parser.parse_args()

    with open(f"paperspace_train_{args.model}_devices_{args.devices}.sh", "w") as f:
        if args.model == "all":
            models = ["schnet", "transformer", "cgcnn", "resnet", "unet", "fno"]
            for model in models:
                generate_one(args, model, f, args.devices)
        else:
            generate_one(args, args.model, f, args.devices)

experiments/gen_train_script.py
lbai-lab/push
import numpy as np
import torch
from timeit import default_timer
from tqdm import tqdm
from typing import *
import wandb

from push.lib.utils import detach_to_cpu, to_device

import sys
sys.path.append('../')
from train_util import get_model, wandb_init, MyTimer


# =============================================================================
# Deep Ensemble Training
# =============================================================================

def train_deep_ensemble(train_loader,args, loss_fn):
    wandb_init(args, train_loader)
    print("Starting Deep Ensemble Training...")
    print(f'Epochs = {args.epochs}, learning rate = {args.learning_rate}, scheduler step = {args.scheduler_step}, scheduler gamma = {args.scheduler_gamma}')
    
    # Create models
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    networks = [get_model(args).to(device) for _ in range(args.num_particles)]

    # Create optimizers and schedulers for each model
    optimizers = [torch.optim.Adam(net.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) for net in networks]
    schedulers = [torch.optim.lr_scheduler.StepLR(opt, step_size=args.scheduler_step, gamma=args.scheduler_gamma) for opt in optimizers]
    
    total_time = 0
    
    # Main training loop
    tq = tqdm(range(args.epochs))
    for e in tq:
        losses = []

        with MyTimer() as my_timer:
            for data, label in train_loader:
                preds = []
                for i, model in enumerate(networks):
                    pred = model(to_device(device, data))
                    preds.append(pred)

                    loss = loss_fn(pred, to_device(device, label))
                    optimizers[i].zero_grad()
                    loss.backward()
                    optimizers[i].step()
                    
                    losses.append(loss.to("cpu").item())

                if args.model == "unet":
                    for scheduler in schedulers:
                        scheduler.step()
        total_time += my_timer.time
        epoch_loss = torch.mean(torch.tensor(losses))
        print("Average Epoch loss", epoch_loss.item())
        
        if args.wandb:
            wandb.log({
                "loss": epoch_loss.item(),
                "time": my_timer.time
            })
            
    average_time = total_time / args.epochs
    print("Average time per epoch:", average_time)
    return networks  # Return the ensemble of models


def test_deep_ensemble(test_loader, args, loss_fn, networks):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Starting Deep Ensemble Testing...")
        
    ensemble_losses = []
    ensemble_uncertainties = []
    total_time = 0
    tq = tqdm(test_loader)
    
    for data, label in tq:
        ensemble_pred = []
        t1 = default_timer()
        
        for model in networks:
            model.eval()  # Set the model to evaluation mode
            # with torch.no_grad():
            pred = model(to_device(device, data))
            ensemble_pred.append(pred)
        
        # Average the predictions 
        ensemble_pred = torch.stack(ensemble_pred)
        avg_pred = torch.mean(ensemble_pred, dim=0)
        
        # Measure uncertainty 
        uncertainty = torch.var(ensemble_pred, dim=0)
        ensemble_uncertainties.append(torch.mean(uncertainty).item())

        loss = loss_fn(avg_pred, to_device(device, label))
        ensemble_losses.append(loss.item())
        
        t2 = default_timer()
        total_time += t2 - t1
        
        if args.wandb:
            wandb.log({
                "test_loss": loss.item(),
                "test_uncertainty": torch.mean(uncertainty).item()
            })
        
    avg_test_loss = np.mean(ensemble_losses)
    avg_test_uncertainty = np.mean(ensemble_uncertainties)
    print(f"Average Test Loss: {avg_test_loss}")
    print(f"Average Test Uncertainty: {avg_test_uncertainty}")
    print(f"Total Test Time: {total_time}")
    
    return avg_test_loss

experiments/infer/ensemble.py
lbai-lab/push
import torch
from tqdm import tqdm
from typing import *
import wandb

from push.lib.utils import detach_to_cpu, to_device

import sys
sys.path.append('../')
from train_util import get_model, wandb_init, MyTimer


# =============================================================================
# Multi Swag Training
# =============================================================================

def update_theta(theta, theta2, tt, tt2, n):
    for param in theta.keys():
        theta[param] = (theta[param]*n+tt[param])/(n+1)
        theta2[param] = (theta2[param]*n+tt2[param])/(n+1)

    return theta, theta2


def train_mswag(train_loader, args, loss_fn):
    wandb_init(args, train_loader)

    print("Training Multi-Swag ...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    networks = [get_model(args).to(device) for _ in range(args.num_particles)]

    # Create optimizers and schedulers for each model
    optimizers = [torch.optim.Adam(net.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) for net in networks]
    schedulers = [torch.optim.lr_scheduler.StepLR(opt, step_size=args.scheduler_step, gamma=args.scheduler_gamma) for opt in optimizers]

    # Pretrain epochs
    total_time = 0
    for e in tqdm(range(args.pretrain_epochs), desc="Epochs"):
        losses = []
        with MyTimer() as my_timer:
            for data, label in train_loader:
                preds = []
                for i, model in enumerate(networks):
                    pred = model(to_device(device, data))
                    preds.append(pred)
                    loss = loss_fn(pred, to_device(device, label))

                    optimizers[i].zero_grad()
                    loss.backward()
                    optimizers[i].step()
                    losses.append(loss.to("cpu").item())
                
                if args.model == "unet":
                    for scheduler in schedulers:
                        scheduler.step()
        epoch_time = my_timer.time
        total_time += epoch_time
        epoch_loss = torch.mean(torch.tensor(losses))
        if args.wandb:
            wandb.log({
                    "pretrain_epoch_time": epoch_time,
                    "pretrain_loss": epoch_loss,
                })        
    
    mom1s = {i: {name: param for name, param in networks[i].named_parameters()} for i in range(args.num_particles)}
    mom2s = {i: {name: param*param for name, param in networks[i].named_parameters()} for i in range(args.num_particles)}
    optimizers = [torch.optim.Adam(nn.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) for nn in networks]

    # SWAG epochs
    for e in tqdm(range(args.swag_epochs), desc="Epochs"):
        with MyTimer() as my_timer:
            losses = []
            for data, label in train_loader:
                for idx, (nn, optimizer) in enumerate(zip(networks, optimizers)):
                    pred = nn(to_device(device, data))
                    loss = loss_fn(pred, to_device(device, label))
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    losses.append(loss.to("cpu").item())
                    nn.zero_grad()
                            
            for idx, nn in enumerate(networks):
                params = {name: param for name, param in nn.named_parameters()}
                params_sq = {name: param*param for name, param in nn.named_parameters()}
                update_theta(mom1s[idx], mom2s[idx], params, params_sq, e+1)
        total_time += my_timer.time
        
        if args.wandb:
            wandb.log({
                "swag_epoch_time": my_timer.time,
                "pretrain_loss": torch.mean(torch.tensor(losses)).item(),
            })
    average_time = total_time / args.epochs
    print("Average time per epoch:", average_time)

    return [mom1s, mom2s]

experiments/infer/mswag.py
lbai-lab/push
import torch
from tqdm import tqdm
import wandb
from typing import *

from push.particle import Particle

import sys
sys.path.append('../')
from train_util import wandb_init, MyTimer


# =============================================================================
# Push Ensemble Training
# =============================================================================

def mk_optim(lr, weight_decay, params):
    return torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)


def _ensemble_main_instrumented(particle: Particle, dataloader, loss_fn, epochs) -> None:
    state = particle.state
    wandb_init(state["args"], dataloader)

    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))
    # Training loop
    for e in tqdm(range(epochs)):
        losses = []
        with MyTimer() as my_timer:
            for data, label in dataloader:
                fut = particle.step(loss_fn, data, label)
                futs = [particle.send(pid, "ENSEMBLE_STEP", loss_fn, data, label) for pid in other_particles]
                losses += [fut.wait()]
                [f.wait() for f in futs]

        # print(f"Average loss {particle.pid}", torch.mean(torch.tensor(losses)))
        if state["args"].wandb:
            wandb.log({
                "time": my_timer.time,
                "leader_loss:": torch.mean(torch.tensor(losses))
            })

experiments/infer/push_ensemble.py
lbai-lab/push
from torch.utils.data import DataLoader
import torch
from tqdm import tqdm
from typing import *
import wandb

import push.bayes.swag as swag
from push.particle import Particle
from push.bayes.utils import flatten, unflatten_like
from push.lib.utils import detach_to_cpu


import sys
sys.path.append('../')
from train_util import wandb_init, MyTimer


# =============================================================================
# PusH Multi-Swag Training
# =============================================================================

def _mswag_particle_instrumented(particle: Particle,
                                 dataloader,
                                 loss_fn: Callable,
                                 pretrain_epochs: int,
                                 swag_epochs: int, swag_pids: list[int]) -> None:
    state = particle.state
    wandb_init(state["args"], dataloader)
    other_pids = [pid for pid in swag_pids if pid != particle.pid]
    
    # Pre-training loop
    for e in tqdm(range(pretrain_epochs)):
        losses = []
        with MyTimer() as my_timer:
            for data, label in dataloader:
                fut = particle.step(loss_fn, data, label)
                futs = [particle.send(pid, "SWAG_STEP", loss_fn, data, label) for pid in other_pids]
                losses += [fut.wait()]
                [f.wait() for f in futs]
        if state["args"].wandb:
            wandb.log({
                "pretrain_epoch_time": my_timer.time,
                "pretrain_loss:": torch.mean(torch.tensor(losses))
            })

    # Initialize swag
    with MyTimer() as my_timer:
        futs = [particle.send(pid, "SWAG_SWAG", True) for pid in other_pids]
        swag._swag_swag(particle, True)
        [f.wait() for f in futs]
    if state["args"].wandb:
        wandb.log({
            "send_time": my_timer.time
        })

    # Swag epochs
    for e in tqdm(range(swag_epochs)):
        losses = []
        with MyTimer() as my_timer:
            for data, label in dataloader:
                # Update
                fut = particle.step(loss_fn, data, label)
                futs = [particle.send(pid, "SWAG_STEP", loss_fn, data, label) for pid in other_pids]
                [f.wait() for f in futs]
                losses += [fut.wait()]
            swag._swag_swag(particle, False)
            futs = [particle.send(pid, "SWAG_SWAG", False) for pid in other_pids]
            [f.wait() for f in futs]
        if state["args"].wandb:
            wandb.log({
                "swag_epoch_time": my_timer.time,
                "leader_loss:": torch.mean(torch.tensor(losses))
            })


# =============================================================================
# PusH Multi-Swag Inference
# =============================================================================

def _mswag_sample_entry_instrumented(particle: Particle,
                                     dataloader: DataLoader,
                                     loss_fn: Callable,
                                     scale: float,
                                     var_clamp: float,
                                     num_samples: int,
                                     num_models) -> None:
    # Unpack state
    state = particle.state
    pid = particle.pid
    other_pids = list(range(1, num_models))

    # Perform SWAG sampling for every particle
    my_ans = _mswag_sample_instrumented(particle, dataloader, loss_fn, scale, var_clamp, num_samples)
    futs = [particle.send(pid, "SWAG_SAMPLE", dataloader, loss_fn, scale, var_clamp, num_samples) for pid in other_pids]
    ans = [f.wait() for f in futs]
    
    # Majority vote of majority prediction
    classes = {k: [0 for i in range(10)] for k in range(10)}
    max_preds = [my_ans['max_preds']] + [ans[i-1]['max_preds'] for i in range(1, num_models)]
    for idx, (data, label) in enumerate(dataloader):
        max_pred = torch.mode(torch.stack([max_preds[i][idx] for i in range(num_models)]), dim=0).values
        for x, y in zip(max_pred, label):
            classes[y.item()][x.item()] += 1
    
    # Majority vote across all SWAG samples
    all_preds = [my_ans['preds']] + [ans[i-1]['preds'] for i in range(1, num_models)]
    classes2 = {k: [0 for i in range(10)] for k in range(10)}
    for idx, (data, label) in enumerate(dataloader):
        preds = []
        for m in range(num_models):
            preds += [torch.stack([all_preds[m][i][idx] for i in range(num_samples)])]
        max_pred = torch.mode(torch.cat(preds, dim=0), dim=0).values
        for x, y in zip(max_pred, label):
            classes2[y.item()][x.item()] += 1

    if state["args"].wandb:
        wandb.log({
            f"max_dist{pid}": str(classes),
            f"flat_max_dist{pid}": str(classes2),
        })


def _mswag_sample_instrumented(particle: Particle,
                               dataloader: DataLoader,
                               loss_fn: Callable,
                               scale: float,
                               var_clamp: float,
                               num_samples: int) -> None:
    """Inspired by https://github.com/wjmaddox/swa_gaussian/blob/master/swag/posteriors/swag.py
    """
    state = particle.state
    pid = particle.pid
    # Gather
    mean_list = [param for param in particle.state[pid]["mom1"]]
    sq_mean_list = [param for param in particle.state[pid]["mom2"]]

    scale_sqrt = scale ** 0.5
    mean = flatten(mean_list)
    sq_mean = flatten(sq_mean_list)

    # Compute original loss
    classes = {k: [0 for i in range(10)] for k in range(10)}
    losses = []
    for data, label in tqdm(dataloader):
        pred = detach_to_cpu(particle.forward(data).wait())
        loss = loss_fn(pred, label)
        cls = pred.softmax(dim=1).argmax(dim=1)
        for x, y in zip(cls, label):
            classes[y.item()][x.item()] += 1
        losses += [loss.detach().to("cpu")]

    if state["args"].wandb:
        wandb.log({
            f"orig_loss{pid}": torch.mean(torch.tensor(losses)),
            f"orig_dist{pid}": str(classes)
        })

    # Compute and store prediction for each SWAG sample
    preds = {i: [] for i in range(num_samples)}
    swag_losses = {}
    for i in range(num_samples):
        # Draw diagonal variance sample
        var = torch.clamp(sq_mean - mean ** 2, var_clamp)
        var_sample = var.sqrt() * torch.randn_like(var, requires_grad=False)

        rand_sample = var_sample

        # Update sample with mean and scale
        sample = mean + scale_sqrt * rand_sample
        sample = sample.unsqueeze(0)

        # Update
        samples_list = unflatten_like(sample, mean_list)

        for param, sample in zip(particle.module.parameters(), samples_list):
            param.data = sample

        swag_losses = []
        for data, label in tqdm(dataloader):
            pred = detach_to_cpu(particle.forward(data).wait())
            cls = pred.softmax(dim=1).argmax(dim=1)
            preds[i] += [cls]
            loss = loss_fn(pred, label)
            swag_losses += [loss.detach().to("cpu")]    

        if state["args"].wandb:
            wandb.log({
                f"swag_loss{pid}": torch.mean(torch.tensor(swag_losses)),
                f"classes_dist{pid}": str(classes)
            })
    
    # Majority prediction for the current SWAG particle
    max_preds = []
    for n in range(len(preds[0])):
        max_preds += [torch.mode(torch.stack([preds[i][n] for i in range(num_samples)]), dim=0).values]

    return {
        "classes": classes, 
        "losses": torch.mean(torch.tensor(losses)),
        "preds": preds,
        "max_preds": max_preds,
    }

experiments/infer/push_mswag.py
lbai-lab/push
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
import wandb
from typing import *

import push.bayes.stein_vgd as svgd
from push.particle import Particle
from push.bayes.utils import flatten
from push.lib.utils import detach_to_cpu

import sys
sys.path.append('../')
from train_util import wandb_init, MyTimer


# =============================================================================
# PusH Stein VGD Training
# =============================================================================

def _svgd_leader_instrumented(particle: Particle,
                              prior, loss_fn: Callable, lengthscale, lr,
                              dataloader: DataLoader, epochs) -> None:
    state = particle.state
    wandb_init(state["args"], dataloader)

    n = len(particle.particle_ids())
    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))

    for e in tqdm(range(epochs)):
        losses = []
        epoch_time = 0
        for data, label in dataloader:
            with MyTimer() as my_timer:
                # 1. Step every particle       
                fut = particle.step(loss_fn, data, label)
                futs = [particle.send(pid, "SVGD_STEP", loss_fn, data, label) for pid in other_particles]
                fut.wait(); [f.wait() for f in futs]

                # 2. Gather every other particles parameters
                particles = {pid: (particle.get(pid) if pid != particle.pid else 
                    list(particle.module.parameters())) for pid in particle.particle_ids()}
                for pid in other_particles:
                    particles[pid] = particles[pid].wait()

                # 3. Compute kernel and kernel gradients
                update = {}
                for pid1, params1 in particles.items():
                    params1 = list(particles[pid1].view().parameters()) if pid1 != particle.pid else params1
                    p_i = flatten(params1)
                    update[pid1] = torch.zeros_like(p_i)
                    for pid2, params2 in particles.items():
                        params2 = list(particles[pid2].view().parameters()) if pid2 != particle.pid else params2

                        # Compute kernel
                        p_j = flatten(params2)
                        p_j_grad = flatten([p.grad if p.grad is not None else torch.zeros_like(p) for p in params2])
                        update[pid1] += svgd.torch_squared_exp_kernel(p_j, p_i, lengthscale) * p_j_grad
                        update[pid1] += svgd.torch_squared_exp_kernel_grad(p_j, p_i, lengthscale).real
                    update[pid1] = update[pid1] / n

                # 4. Send kernel
                futs = [particle.send(pid, "SVGD_FOLLOW", lr, update[pid]) for pid in other_particles]
                [f.wait() for f in futs]
                svgd._svgd_follow(particle, lr, update[particle.pid])
            epoch_time += my_timer.time


            losses += [loss_fn(detach_to_cpu(particle.forward(data).wait()), label)]
        if state["args"].wandb:
            wandb.log({
                "time": epoch_time,
                "leader_loss:": torch.mean(torch.tensor(losses))
            })


def _svgd_leader_instrumented_memeff(particle: Particle, prior, loss_fn: Callable, lengthscale, lr, dataloader: DataLoader, epochs) -> None:
    print("DOING 2")
    state = particle.state
    wandb_init(state["args"], dataloader)

    n = len(particle.particle_ids())
    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))

    for e in tqdm(range(epochs)):
        losses = []
        epoch_time = 0
        for data, label in tqdm(dataloader,desc="data_loader"):
            with MyTimer() as my_timer:
                
                fut = particle.step(loss_fn, data, label)
                futs = [particle.send(pid, "SVGD_STEP", loss_fn, data, label) for pid in other_particles]
                fut.wait(); [f.wait() for f in futs]

                # 2. Gather every other particles parameters
                particles = {pid: (particle.get(pid) if pid != particle.pid else 
                    list(particle.module.parameters())) for pid in particle.particle_ids()}
                for pid in other_particles:
                    particles[pid] = particles[pid].wait()

                def compute_update(pid1, params1):
                    params1 = list(particles[pid1].view().parameters()) if pid1 != particle.pid else params1
                    p_i = flatten(params1)
                    update = torch.zeros_like(p_i)
                    for pid2, params2 in particles.items():
                        params2 = list(particles[pid2].view().parameters()) if pid2 != particle.pid else params2

                        # Compute kernel
                        p_j = flatten(params2)
                        p_j_grad = flatten([p.grad if p.grad is not None else torch.zeros_like(p) for p in params2])
                        diff = (p_j - p_i) / lengthscale
                        radius2 = torch.dot(diff, diff)
                        k = torch.exp(-0.5 * radius2).item()
                        diff.mul_(-k / lengthscale)
                        update.add_(p_j_grad, alpha=k.real)
                        update.add_(diff.real)
                        # update += svgd.torch_squared_exp_kernel(p_j, p_i, lengthscale).real * p_j_grad
                        # update += svgd.torch_squared_exp_kernel_grad(p_j, p_i, lengthscale).real
                    update = update / n
                    return update

                # 3. Compute kernel and kernel gradients
                for pid1, params1 in particles.items():
                    # 4. Send kernel
                    if pid1 != particle.pid:
                        update = compute_update(pid1, params1)
                        particle.send(pid, "SVGD_FOLLOW", lr, update).wait()
                update = compute_update(particle.pid, particles[particle.pid])
                svgd._svgd_follow(particle, lr, update)
            epoch_time += my_timer.time
            # Record losses
            loss= loss_fn(detach_to_cpu(particle.forward(data).wait()), label) 
            losses += [loss]
        if state["args"].wandb:
            wandb.log({
                "time": epoch_time,
                "leader_loss:": torch.mean(torch.tensor(losses))
            })

experiments/infer/push_stein_vgd.py
lbai-lab/push
import torch
from timeit import default_timer
from tqdm import tqdm
from typing import *
import wandb

from push.lib.utils import detach_to_cpu, to_device

import sys
sys.path.append('../')
from train_util import get_model



# =============================================================================
# Standard Training
# =============================================================================

def train_standard(train_loader, args, loss_fn):
    print("Starting Standard Training...")
    print(f'Epochs = {args.epochs}, learning rate = {args.learning_rate}, scheduler step = {args.scheduler_step}, scheduler gamma = {args.scheduler_gamma}')
    # Create model
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = get_model(args).to(device)

    # Create optimizers and schedulers for each particle
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.scheduler_step, gamma=args.scheduler_gamma)

    total_time = 0
    tq = tqdm(range(args.epochs))
    for e in tq:
        losses = []
        t1 = default_timer()
        for data, label in train_loader:
            pred = model(to_device(device, data))
            loss = loss_fn(pred, to_device(device, label))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            losses += [loss.to("cpu")]
        
        if args.model == "unet":
            scheduler.step()
        
        t2 = default_timer()
        total_time += t2 - t1
        epoch_loss = torch.mean(torch.tensor(losses))
        print("Average Epoch loss", epoch_loss)

        if args.wandb:
            wandb.log({
                    "loss": torch.mean(torch.tensor(losses)).item(),
                    "time": t2 - t1})
            
    average_time = total_time / args.epochs
    print("Average time per epoch:", average_time)
    return model


def test_standard(test_loader, args, loss_fn, trained_model):
    print("Testing Standard Training...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    test_losses = []
    total_batches = 0
    for data, label in tqdm(test_loader, desc="testloader"):
        pred = trained_model(to_device(device, data))
        test_loss = loss_fn(pred, to_device(device, label))
        test_losses.append(detach_to_cpu(test_loss).item())
        total_batches += 1

        if args.wandb:
            wandb.log({"test loss " + str(args.dataset): test_loss.item()})
    # Compute average test loss
    average_test_loss = sum(test_losses) / total_batches
    print(f"Average Test Loss: {average_test_loss}")

experiments/infer/standard.py
lbai-lab/push
import argparse
from datetime import datetime
import pytz
import torch
from timeit import default_timer
import time
from tqdm import tqdm
from typing import *
import sys
sys.path.append('../')
import wandb

from push.bayes.utils import flatten, unflatten_like
from push.bayes.stein_vgd import torch_squared_exp_kernel, torch_squared_exp_kernel_grad
from push.lib.utils import detach_to_cpu, to_device
from train_util import get_model, wandb_init, MyTimer


# =============================================================================
# SVGD Training
# =============================================================================

def svgd_step_precompute( n, networks, bandwidth, lr): #parameterwise version
    network_params = [[] for _ in range(n)]

    for i in range(n):
        network_params[i] = [param for param in networks[i].parameters()]

    K = {i: torch.zeros(n, device="cuda") for i in range(n)}
    K_grad = {i: [] for i in range(n)}
    
    for i in range(n):
        params1 = list(network_params[i]) 
        p_i = flatten(params1)
        for j in range(n):
            update = torch.zeros_like(p_i)
            params2 = list(network_params[j]) 
            # Compute kernel
            p_j = flatten(params2)
            diff = (p_j - p_i) / bandwidth
            radius2 = torch.dot(diff, diff)
            k = torch.exp(-0.5 * radius2).item()
            diff.mul_(-k / bandwidth)
            update.add_(diff.real)
            K[i][j] = k
            K_grad[i] += [update]
            # K[i][j] = torch_squared_exp_kernel(p_i, p_j, bandwidth).real
            # K_grad[i] += [torch_squared_exp_kernel_grad(p_i, p_j, bandwidth).real]    
   
    for i in range(n):
        acc = [torch.zeros_like(p) for p in network_params[i]]
        for j in range(n):

            pj = list(network_params[j])
            k_grads_j = unflatten_like(K_grad[j][i].unsqueeze(0),pj) 
            
            for idx, tmp in enumerate(acc): 
                #print(K[i][j].dtype, k_grads_j[idx].dtype, pj[idx].grad.dtype, tmp.dtype)
                if network_params[j][idx].grad is not None:
                    tmp.add_(pj[idx].grad, alpha=K[j][i].item()/n)
                tmp.add_(k_grads_j[idx].real, alpha=1/n) 
        with torch.no_grad():
            for idx, p in enumerate(network_params[i]):
                if p.grad is not None:
                    p.add_(acc[idx], alpha=-lr)


def train_svgd(train_loader, args, loss_fn):
    wandb_init(args, train_loader)

    print("Training SVGD.....")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    networks = [get_model(args).to(device) for _ in range(args.num_particles)]
    total_time = 0
    for e in tqdm(range(args.epochs), desc="Epochs"):
        losses = []
        with MyTimer() as my_timer:
            for data, label in train_loader:
                # Gradients
                for network in networks:
                    network.zero_grad()
                    pred = network(to_device(device, data))
                    loss = loss_fn(pred, to_device(device, label))
                    losses += [loss.to("cpu").item()]
                    loss.backward()  

                svgd_step_precompute( args.num_particles, networks, args.bandwidth, args.learning_rate)
        epoch_time = my_timer.time
        total_time += epoch_time
        print("Time per svgd epoch:", epoch_time)
        print("Average loss:", torch.mean(torch.tensor(losses)))
        if args.wandb:
            wandb.log({
                "epoch_loss": torch.mean(torch.tensor(losses)).item(),
                "time": epoch_time,
            })
    if args.wandb:
        wandb.log({
            "average time": total_time/args.epochs,         
        })
    return networks


def test_svgd(test_loader, args, loss_fn, networks):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Testing test_loss")
    test_losses = []
    
    for data, label in test_loader:
        preds = []
        
        #Average predictions 
        for network in networks:
            preds.append(network(to_device(device, data)))
        preds_stacked = torch.stack(preds, dim=0)
        pred_mean = torch.mean(preds_stacked, dim=0)

        test_loss = loss_fn(pred_mean, to_device(device, label))
        if args.wandb:
            wandb.log({"test loss " + args.dataset: test_loss})
        test_losses += [test_loss.to("cpu")]
    avg_test_loss = torch.mean(torch.tensor(test_losses))
    print("Average Test Loss: Standard", avg_test_loss)
    if args.wandb:
        wandb.log({"average test loss " + args.dataset: avg_test_loss})

experiments/infer/stein_vgd.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import os
import sys
import logging
import torch
import torch.nn as nn

proj_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
if proj_root not in sys.path: sys.path.insert(1, proj_root)
from torch_geometric.nn import radius_graph
from nns.cgcnn.base_util import (compute_neighbors, conditional_grad, get_pbc_distances, radius_graph_pbc,)


class BaseModel(nn.Module):
    def __init__(self, num_atoms=None, bond_feat_dim=None, num_targets=None):
        super(BaseModel, self).__init__()
        self.num_atoms = num_atoms
        self.bond_feat_dim = bond_feat_dim
        self.num_targets = num_targets

    def forward(self, data):
        raise NotImplementedError

    def generate_graph(
        self,
        data,
        cutoff=None,
        max_neighbors=None,
        use_pbc=None,
        otf_graph=None,
    ):
        cutoff = cutoff or self.cutoff
        max_neighbors = max_neighbors or self.max_neighbors
        use_pbc = use_pbc or self.use_pbc
        otf_graph = otf_graph or self.otf_graph

        if not otf_graph:
            try:
                edge_index = data.edge_index

                if use_pbc:
                    cell_offsets = data.cell_offsets
                    neighbors = data.neighbors

            except AttributeError:
                logging.warning(
                    "Turning otf_graph=True as required attributes not present in data object"
                )
                otf_graph = True

        if use_pbc:
            if otf_graph:
                edge_index, cell_offsets, neighbors = radius_graph_pbc(data, cutoff, max_neighbors)

            out = get_pbc_distances(
                data.pos,
                edge_index,
                data.cell,
                cell_offsets,
                neighbors,
                return_offsets=True,
                return_distance_vec=True,
            )

            edge_index = out["edge_index"]
            edge_dist = out["distances"]
            cell_offset_distances = out["offsets"]
            distance_vec = out["distance_vec"]
        else:
            if otf_graph:
                edge_index = radius_graph(data.pos, r=cutoff, batch=data.batch, max_num_neighbors=max_neighbors,)

            j, i = edge_index
            distance_vec = data.pos[j] - data.pos[i]

            edge_dist = distance_vec.norm(dim=-1)
            cell_offsets = torch.zeros(edge_index.shape[1], 3, device=data.pos.device)
            cell_offset_distances = torch.zeros_like(cell_offsets, device=data.pos.device)
            neighbors = compute_neighbors(data, edge_index)

        return (
            edge_index,
            edge_dist,
            distance_vec,
            cell_offsets,
            cell_offset_distances,
            neighbors,)

    #@property
    #def num_params(self):
    #    return sum(p.numel() for p in self.parameters())
experiments/nns/cgcnn/base.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

import ast
import collections
import copy
import importlib
import itertools
import json
import logging
import os
import sys
import time
from argparse import Namespace
from bisect import bisect
from contextlib import contextmanager
from dataclasses import dataclass
from functools import wraps
from itertools import product
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional

import numpy as np
import torch
import torch.nn as nn
import torch_geometric
import yaml
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from torch_geometric.data import Data
from torch_geometric.utils import remove_self_loops
from torch_scatter import segment_coo, segment_csr

if TYPE_CHECKING:
    from torch.nn.modules.module import _IncompatibleKeys


def pyg2_data_transform(data: Data):
    """
    if we're on the new pyg (2.0 or later) and if the Data stored is in older format
    we need to convert the data to the new format
    """
    if torch_geometric.__version__ >= "2.0" and "_store" not in data.__dict__:
        return Data(
            **{k: v for k, v in data.__dict__.items() if v is not None}
        )

    return data


def save_checkpoint(
    state, checkpoint_dir="checkpoints/", checkpoint_file="checkpoint.pt"
):
    filename = os.path.join(checkpoint_dir, checkpoint_file)
    torch.save(state, filename)
    return filename


class Complete(object):
    def __call__(self, data):
        device = data.edge_index.device

        row = torch.arange(data.num_nodes, dtype=torch.long, device=device)
        col = torch.arange(data.num_nodes, dtype=torch.long, device=device)

        row = row.view(-1, 1).repeat(1, data.num_nodes).view(-1)
        col = col.repeat(data.num_nodes)
        edge_index = torch.stack([row, col], dim=0)

        edge_attr = None
        if data.edge_attr is not None:
            idx = data.edge_index[0] * data.num_nodes + data.edge_index[1]
            size = list(data.edge_attr.size())
            size[0] = data.num_nodes * data.num_nodes
            edge_attr = data.edge_attr.new_zeros(size)
            edge_attr[idx] = data.edge_attr

        edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
        data.edge_attr = edge_attr
        data.edge_index = edge_index

        return data


def warmup_lr_lambda(current_step, optim_config):
    """Returns a learning rate multiplier.
    Till `warmup_steps`, learning rate linearly increases to `initial_lr`,
    and then gets multiplied by `lr_gamma` every time a milestone is crossed.
    """

    # keep this block for older configs that have warmup_epochs instead of warmup_steps
    # and lr_milestones are defined in epochs
    if (
        any(x < 100 for x in optim_config["lr_milestones"])
        or "warmup_epochs" in optim_config
    ):
        raise Exception(
            "ConfigError: please define lr_milestones in steps not epochs and define warmup_steps instead of warmup_epochs"
        )

    if current_step <= optim_config["warmup_steps"]:
        alpha = current_step / float(optim_config["warmup_steps"])
        return optim_config["warmup_factor"] * (1.0 - alpha) + alpha
    else:
        idx = bisect(optim_config["lr_milestones"], current_step)
        return pow(optim_config["lr_gamma"], idx)


def print_cuda_usage():
    print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 * 1024))
    print(
        "Max Memory Allocated:",
        torch.cuda.max_memory_allocated() / (1024 * 1024),
    )
    print("Memory Cached:", torch.cuda.memory_cached() / (1024 * 1024))
    print("Max Memory Cached:", torch.cuda.max_memory_cached() / (1024 * 1024))


def conditional_grad(dec):
    "Decorator to enable/disable grad depending on whether force/energy predictions are being made"
    # Adapted from https://stackoverflow.com/questions/60907323/accessing-class-property-as-decorator-argument
    def decorator(func):
        @wraps(func)
        def cls_method(self, *args, **kwargs):
            f = func
            if self.regress_forces and not getattr(self, "direct_forces", 0):
                f = dec(func)
            return f(self, *args, **kwargs)

        return cls_method

    return decorator


def plot_histogram(data, xlabel="", ylabel="", title=""):
    assert isinstance(data, list)

    # Preset
    fig = Figure(figsize=(5, 4), dpi=150)
    canvas = FigureCanvas(fig)
    ax = fig.gca()

    # Plot
    ax.hist(data, bins=20, rwidth=0.9, zorder=3)

    # Axes
    ax.grid(color="0.95", zorder=0)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    fig.tight_layout(pad=2)

    # Return numpy array
    canvas.draw()
    image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    image_from_plot = image_from_plot.reshape(
        fig.canvas.get_width_height()[::-1] + (3,)
    )

    return image_from_plot


# Override the collation method in `pytorch_geometric.data.InMemoryDataset`
def collate(data_list):
    keys = data_list[0].keys
    data = data_list[0].__class__()

    for key in keys:
        data[key] = []
    slices = {key: [0] for key in keys}

    for item, key in product(data_list, keys):
        data[key].append(item[key])
        if torch.is_tensor(item[key]):
            s = slices[key][-1] + item[key].size(
                item.__cat_dim__(key, item[key])
            )
        elif isinstance(item[key], int) or isinstance(item[key], float):
            s = slices[key][-1] + 1
        else:
            raise ValueError("Unsupported attribute type")
        slices[key].append(s)

    if hasattr(data_list[0], "__num_nodes__"):
        data.__num_nodes__ = []
        for item in data_list:
            data.__num_nodes__.append(item.num_nodes)

    for key in keys:
        if torch.is_tensor(data_list[0][key]):
            data[key] = torch.cat(
                data[key], dim=data.__cat_dim__(key, data_list[0][key])
            )
        else:
            data[key] = torch.tensor(data[key])
        slices[key] = torch.tensor(slices[key], dtype=torch.long)

    return data, slices


def add_edge_distance_to_graph(
    batch,
    device="cpu",
    dmin=0.0,
    dmax=6.0,
    num_gaussians=50,
):
    # Make sure x has positions.
    if not all(batch.pos[0][:] == batch.x[0][-3:]):
        batch.x = torch.cat([batch.x, batch.pos.float()], dim=1)
    # First set computations to be tracked for positions.
    batch.x = batch.x.requires_grad_(True)
    # Then compute Euclidean distance between edge endpoints.
    pdist = torch.nn.PairwiseDistance(p=2.0)
    distances = pdist(
        batch.x[batch.edge_index[0]][:, -3:],
        batch.x[batch.edge_index[1]][:, -3:],
    )
    # Expand it using a gaussian basis filter.
    gdf_filter = torch.linspace(dmin, dmax, num_gaussians)
    var = gdf_filter[1] - gdf_filter[0]
    gdf_filter, var = gdf_filter.to(device), var.to(device)
    gdf_distances = torch.exp(
        -((distances.view(-1, 1) - gdf_filter) ** 2) / var**2
    )
    # Reassign edge attributes.
    batch.edge_weight = distances
    batch.edge_attr = gdf_distances.float()
    return batch


def _import_local_file(path: Path, *, project_root: Path):
    """
    Imports a Python file as a module
    :param path: The path to the file to import
    :type path: Path
    :param project_root: The root directory of the project (i.e., the "ocp" folder)
    :type project_root: Path
    """

    path = path.resolve()
    project_root = project_root.resolve()

    module_name = ".".join(
        path.absolute()
        .relative_to(project_root.absolute())
        .with_suffix("")
        .parts
    )
    logging.debug(f"Resolved module name of {path} to {module_name}")
    importlib.import_module(module_name)


def setup_experimental_imports(project_root: Path):
    experimental_folder = (project_root / "experimental").resolve()
    if not experimental_folder.exists() or not experimental_folder.is_dir():
        return

    experimental_files = [
        f.resolve().absolute() for f in experimental_folder.rglob("*.py")
    ]
    # Ignore certain directories within experimental
    ignore_file = experimental_folder / ".ignore"
    if ignore_file.exists():
        with open(ignore_file, "r") as f:
            for line in f.read().splitlines():
                for ignored_file in (experimental_folder / line).rglob("*.py"):
                    experimental_files.remove(
                        ignored_file.resolve().absolute()
                    )

    for f in experimental_files:
        _import_local_file(f, project_root=project_root)


def _get_project_root():
    """
    Gets the root folder of the project (the "ocp" folder)
    :return: The absolute path to the project root.
    """
    from ocpmodels.common.registry import registry

    # Automatically load all of the modules, so that
    # they register with registry
    root_folder = registry.get("ocpmodels_root", no_warning=True)

    if root_folder is not None:
        assert isinstance(root_folder, str), "ocpmodels_root must be a string"
        root_folder = Path(root_folder).resolve().absolute()
        assert root_folder.exists(), f"{root_folder} does not exist"
        assert root_folder.is_dir(), f"{root_folder} is not a directory"
    else:
        root_folder = Path(__file__).resolve().absolute().parent.parent

    # root_folder is the "ocpmodes" folder, so we need to go up one more level
    return root_folder.parent


# Copied from https://github.com/facebookresearch/mmf/blob/master/mmf/utils/env.py#L89.
def setup_imports(config: Optional[dict] = None):
    from ocpmodels.common.registry import registry

    skip_experimental_imports = (config or {}).get(
        "skip_experimental_imports", None
    )

    # First, check if imports are already setup
    has_already_setup = registry.get("imports_setup", no_warning=True)
    if has_already_setup:
        return

    try:
        project_root = _get_project_root()
        logging.info(f"Project root: {project_root}")
        importlib.import_module("ocpmodels.common.logger")

        import_keys = ["trainers", "datasets", "models", "tasks"]
        for key in import_keys:
            for f in (project_root / "ocpmodels" / key).rglob("*.py"):
                _import_local_file(f, project_root=project_root)

        if not skip_experimental_imports:
            setup_experimental_imports(project_root)
    finally:
        registry.register("imports_setup", True)


def dict_set_recursively(dictionary, key_sequence, val):
    top_key = key_sequence.pop(0)
    if len(key_sequence) == 0:
        dictionary[top_key] = val
    else:
        if top_key not in dictionary:
            dictionary[top_key] = {}
        dict_set_recursively(dictionary[top_key], key_sequence, val)


def parse_value(value):
    """
    Parse string as Python literal if possible and fallback to string.
    """
    try:
        return ast.literal_eval(value)
    except (ValueError, SyntaxError):
        # Use as string if nothing else worked
        return value


def create_dict_from_args(args: list, sep: str = "."):
    """
    Create a (nested) dictionary from console arguments.
    Keys in different dictionary levels are separated by sep.
    """
    return_dict = {}
    for arg in args:
        arg = arg.strip("--")
        keys_concat, val = arg.split("=")
        val = parse_value(val)
        key_sequence = keys_concat.split(sep)
        dict_set_recursively(return_dict, key_sequence, val)
    return return_dict


def load_config(path: str, previous_includes: list = []):
    path = Path(path)
    if path in previous_includes:
        raise ValueError(
            f"Cyclic config include detected. {path} included in sequence {previous_includes}."
        )
    previous_includes = previous_includes + [path]

    direct_config = yaml.safe_load(open(path, "r"))

    # Load config from included files.
    if "includes" in direct_config:
        includes = direct_config.pop("includes")
    else:
        includes = []
    if not isinstance(includes, list):
        raise AttributeError(
            "Includes must be a list, '{}' provided".format(type(includes))
        )

    config = {}
    duplicates_warning = []
    duplicates_error = []

    for include in includes:
        include_config, inc_dup_warning, inc_dup_error = load_config(
            include, previous_includes
        )
        duplicates_warning += inc_dup_warning
        duplicates_error += inc_dup_error

        # Duplicates between includes causes an error
        config, merge_dup_error = merge_dicts(config, include_config)
        duplicates_error += merge_dup_error

    # Duplicates between included and main file causes warnings
    config, merge_dup_warning = merge_dicts(config, direct_config)
    duplicates_warning += merge_dup_warning

    return config, duplicates_warning, duplicates_error


def build_config(args, args_override):
    config, duplicates_warning, duplicates_error = load_config(args.config_yml)
    if len(duplicates_warning) > 0:
        logging.warning(
            f"Overwritten config parameters from included configs "
            f"(non-included parameters take precedence): {duplicates_warning}"
        )
    if len(duplicates_error) > 0:
        raise ValueError(
            f"Conflicting (duplicate) parameters in simultaneously "
            f"included configs: {duplicates_error}"
        )

    # Check for overridden parameters.
    if args_override != []:
        overrides = create_dict_from_args(args_override)
        config, _ = merge_dicts(config, overrides)

    # Some other flags.
    config["mode"] = args.mode
    config["identifier"] = args.identifier
    config["timestamp_id"] = args.timestamp_id
    config["seed"] = args.seed
    config["is_debug"] = args.debug
    config["run_dir"] = args.run_dir
    config["print_every"] = args.print_every
    config["amp"] = args.amp
    config["checkpoint"] = args.checkpoint
    config["cpu"] = args.cpu
    # Submit
    config["submit"] = args.submit
    config["summit"] = args.summit
    # Distributed
    config["local_rank"] = args.local_rank
    config["distributed_port"] = args.distributed_port
    config["world_size"] = args.num_nodes * args.num_gpus
    config["distributed_backend"] = args.distributed_backend
    config["noddp"] = args.no_ddp
    config["gp_gpus"] = args.gp_gpus

    return config


def create_grid(base_config, sweep_file):
    def _flatten_sweeps(sweeps, root_key="", sep="."):
        flat_sweeps = []
        for key, value in sweeps.items():
            new_key = root_key + sep + key if root_key else key
            if isinstance(value, collections.MutableMapping):
                flat_sweeps.extend(_flatten_sweeps(value, new_key).items())
            else:
                flat_sweeps.append((new_key, value))
        return collections.OrderedDict(flat_sweeps)

    def _update_config(config, keys, override_vals, sep="."):
        for key, value in zip(keys, override_vals):
            key_path = key.split(sep)
            child_config = config
            for name in key_path[:-1]:
                child_config = child_config[name]
            child_config[key_path[-1]] = value
        return config

    sweeps = yaml.safe_load(open(sweep_file, "r"))
    flat_sweeps = _flatten_sweeps(sweeps)
    keys = list(flat_sweeps.keys())
    values = list(itertools.product(*flat_sweeps.values()))

    configs = []
    for i, override_vals in enumerate(values):
        config = copy.deepcopy(base_config)
        config = _update_config(config, keys, override_vals)
        config["identifier"] = config["identifier"] + f"_run{i}"
        configs.append(config)
    return configs


def save_experiment_log(args, jobs, configs):
    log_file = args.logdir / "exp" / time.strftime("%Y-%m-%d-%I-%M-%S%p.log")
    log_file.parent.mkdir(exist_ok=True, parents=True)
    with open(log_file, "w") as f:
        for job, config in zip(jobs, configs):
            print(
                json.dumps(
                    {
                        "config": config,
                        "slurm_id": job.job_id,
                        "timestamp": time.strftime("%I:%M:%S%p %Z %b %d, %Y"),
                    }
                ),
                file=f,
            )
    return log_file


def get_pbc_distances(
    pos,
    edge_index,
    cell,
    cell_offsets,
    neighbors,
    return_offsets=False,
    return_distance_vec=False,
):
    row, col = edge_index

    distance_vectors = pos[row] - pos[col]

    # correct for pbc
    neighbors = neighbors.to(cell.device)
    cell = torch.repeat_interleave(cell, neighbors, dim=0)
    offsets = cell_offsets.float().view(-1, 1, 3).bmm(cell.float()).view(-1, 3)
    distance_vectors += offsets

    # compute distances
    distances = distance_vectors.norm(dim=-1)

    # redundancy: remove zero distances
    nonzero_idx = torch.arange(len(distances))[distances != 0]
    edge_index = edge_index[:, nonzero_idx]
    distances = distances[nonzero_idx]

    out = {
        "edge_index": edge_index,
        "distances": distances,
    }

    if return_distance_vec:
        out["distance_vec"] = distance_vectors[nonzero_idx]

    if return_offsets:
        out["offsets"] = offsets[nonzero_idx]

    return out


def radius_graph_pbc(
    data, radius, max_num_neighbors_threshold, pbc=[True, True, False]
):
    device = data.pos.device
    batch_size = len(data.natoms)

    # position of the atoms
    atom_pos = data.pos

    # Before computing the pairwise distances between atoms, first create a list of atom indices to compare for the entire batch
    num_atoms_per_image = data.natoms
    num_atoms_per_image_sqr = (num_atoms_per_image**2).long()

    # index offset between images
    index_offset = (
        torch.cumsum(num_atoms_per_image, dim=0) - num_atoms_per_image
    )

    index_offset_expand = torch.repeat_interleave(
        index_offset, num_atoms_per_image_sqr
    )
    num_atoms_per_image_expand = torch.repeat_interleave(
        num_atoms_per_image, num_atoms_per_image_sqr
    )

    # Compute a tensor containing sequences of numbers that range from 0 to num_atoms_per_image_sqr for each image
    # that is used to compute indices for the pairs of atoms. This is a very convoluted way to implement
    # the following (but 10x faster since it removes the for loop)
    # for batch_idx in range(batch_size):
    #    batch_count = torch.cat([batch_count, torch.arange(num_atoms_per_image_sqr[batch_idx], device=device)], dim=0)
    num_atom_pairs = torch.sum(num_atoms_per_image_sqr)
    index_sqr_offset = (
        torch.cumsum(num_atoms_per_image_sqr, dim=0) - num_atoms_per_image_sqr
    )
    index_sqr_offset = torch.repeat_interleave(
        index_sqr_offset, num_atoms_per_image_sqr
    )
    atom_count_sqr = (
        torch.arange(num_atom_pairs, device=device) - index_sqr_offset
    )

    # Compute the indices for the pairs of atoms (using division and mod)
    # If the systems get too large this apporach could run into numerical precision issues
    index1 = (
        torch.div(
            atom_count_sqr, num_atoms_per_image_expand, rounding_mode="floor"
        )
    ) + index_offset_expand
    index2 = (
        atom_count_sqr % num_atoms_per_image_expand
    ) + index_offset_expand
    # Get the positions for each atom
    pos1 = torch.index_select(atom_pos, 0, index1)
    pos2 = torch.index_select(atom_pos, 0, index2)

    # Calculate required number of unit cells in each direction.
    # Smallest distance between planes separated by a1 is
    # 1 / ||(a2 x a3) / V||_2, since a2 x a3 is the area of the plane.
    # Note that the unit cell volume V = a1 * (a2 x a3) and that
    # (a2 x a3) / V is also the reciprocal primitive vector
    # (crystallographer's definition).

    cross_a2a3 = torch.cross(data.cell[:, 1], data.cell[:, 2], dim=-1)
    cell_vol = torch.sum(data.cell[:, 0] * cross_a2a3, dim=-1, keepdim=True)

    if pbc[0]:
        inv_min_dist_a1 = torch.norm(cross_a2a3 / cell_vol, p=2, dim=-1)
        rep_a1 = torch.ceil(radius * inv_min_dist_a1)
    else:
        rep_a1 = data.cell.new_zeros(1)

    if pbc[1]:
        cross_a3a1 = torch.cross(data.cell[:, 2], data.cell[:, 0], dim=-1)
        inv_min_dist_a2 = torch.norm(cross_a3a1 / cell_vol, p=2, dim=-1)
        rep_a2 = torch.ceil(radius * inv_min_dist_a2)
    else:
        rep_a2 = data.cell.new_zeros(1)

    if pbc[2]:
        cross_a1a2 = torch.cross(data.cell[:, 0], data.cell[:, 1], dim=-1)
        inv_min_dist_a3 = torch.norm(cross_a1a2 / cell_vol, p=2, dim=-1)
        rep_a3 = torch.ceil(radius * inv_min_dist_a3)
    else:
        rep_a3 = data.cell.new_zeros(1)

    # Take the max over all images for uniformity. This is essentially padding.
    # Note that this can significantly increase the number of computed distances
    # if the required repetitions are very different between images
    # (which they usually are). Changing this to sparse (scatter) operations
    # might be worth the effort if this function becomes a bottleneck.
    max_rep = [rep_a1.max(), rep_a2.max(), rep_a3.max()]

    # Tensor of unit cells
    cells_per_dim = [
        torch.arange(-rep, rep + 1, device=device, dtype=torch.float)
        for rep in max_rep
    ]
    unit_cell = torch.cartesian_prod(*cells_per_dim)
    num_cells = len(unit_cell)
    unit_cell_per_atom = unit_cell.view(1, num_cells, 3).repeat(
        len(index2), 1, 1
    )
    unit_cell = torch.transpose(unit_cell, 0, 1)
    unit_cell_batch = unit_cell.view(1, 3, num_cells).expand(
        batch_size, -1, -1
    )

    # Compute the x, y, z positional offsets for each cell in each image
    data_cell = torch.transpose(data.cell, 1, 2)
    pbc_offsets = torch.bmm(data_cell, unit_cell_batch)
    pbc_offsets_per_atom = torch.repeat_interleave(
        pbc_offsets, num_atoms_per_image_sqr, dim=0
    )

    # Expand the positions and indices for the 9 cells
    pos1 = pos1.view(-1, 3, 1).expand(-1, -1, num_cells)
    pos2 = pos2.view(-1, 3, 1).expand(-1, -1, num_cells)
    index1 = index1.view(-1, 1).repeat(1, num_cells).view(-1)
    index2 = index2.view(-1, 1).repeat(1, num_cells).view(-1)
    # Add the PBC offsets for the second atom
    pos2 = pos2 + pbc_offsets_per_atom

    # Compute the squared distance between atoms
    atom_distance_sqr = torch.sum((pos1 - pos2) ** 2, dim=1)
    atom_distance_sqr = atom_distance_sqr.view(-1)

    # Remove pairs that are too far apart
    mask_within_radius = torch.le(atom_distance_sqr, radius * radius)
    # Remove pairs with the same atoms (distance = 0.0)
    mask_not_same = torch.gt(atom_distance_sqr, 0.0001)
    mask = torch.logical_and(mask_within_radius, mask_not_same)
    index1 = torch.masked_select(index1, mask)
    index2 = torch.masked_select(index2, mask)
    unit_cell = torch.masked_select(
        unit_cell_per_atom.view(-1, 3), mask.view(-1, 1).expand(-1, 3)
    )
    unit_cell = unit_cell.view(-1, 3)
    atom_distance_sqr = torch.masked_select(atom_distance_sqr, mask)

    mask_num_neighbors, num_neighbors_image = get_max_neighbors_mask(
        natoms=data.natoms,
        index=index1,
        atom_distance=atom_distance_sqr,
        max_num_neighbors_threshold=max_num_neighbors_threshold,
    )

    if not torch.all(mask_num_neighbors):
        # Mask out the atoms to ensure each atom has at most max_num_neighbors_threshold neighbors
        index1 = torch.masked_select(index1, mask_num_neighbors)
        index2 = torch.masked_select(index2, mask_num_neighbors)
        unit_cell = torch.masked_select(
            unit_cell.view(-1, 3), mask_num_neighbors.view(-1, 1).expand(-1, 3)
        )
        unit_cell = unit_cell.view(-1, 3)

    edge_index = torch.stack((index2, index1))

    return edge_index, unit_cell, num_neighbors_image


def get_max_neighbors_mask(
    natoms, index, atom_distance, max_num_neighbors_threshold
):
    """
    Give a mask that filters out edges so that each atom has at most
    `max_num_neighbors_threshold` neighbors.
    Assumes that `index` is sorted.
    """
    device = natoms.device
    num_atoms = natoms.sum()

    # Get number of neighbors
    # segment_coo assumes sorted index
    ones = index.new_ones(1).expand_as(index)
    num_neighbors = segment_coo(ones, index, dim_size=num_atoms)
    max_num_neighbors = num_neighbors.max()
    num_neighbors_thresholded = num_neighbors.clamp(
        max=max_num_neighbors_threshold
    )

    # Get number of (thresholded) neighbors per image
    image_indptr = torch.zeros(
        natoms.shape[0] + 1, device=device, dtype=torch.long
    )
    image_indptr[1:] = torch.cumsum(natoms, dim=0)
    num_neighbors_image = segment_csr(num_neighbors_thresholded, image_indptr)

    # If max_num_neighbors is below the threshold, return early
    if (
        max_num_neighbors <= max_num_neighbors_threshold
        or max_num_neighbors_threshold <= 0
    ):
        mask_num_neighbors = torch.tensor(
            [True], dtype=bool, device=device
        ).expand_as(index)
        return mask_num_neighbors, num_neighbors_image

    # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors.
    # Fill with infinity so we can easily remove unused distances later.
    distance_sort = torch.full(
        [num_atoms * max_num_neighbors], np.inf, device=device
    )

    # Create an index map to map distances from atom_distance to distance_sort
    # index_sort_map assumes index to be sorted
    index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors
    index_neighbor_offset_expand = torch.repeat_interleave(
        index_neighbor_offset, num_neighbors
    )
    index_sort_map = (
        index * max_num_neighbors
        + torch.arange(len(index), device=device)
        - index_neighbor_offset_expand
    )
    distance_sort.index_copy_(0, index_sort_map, atom_distance)
    distance_sort = distance_sort.view(num_atoms, max_num_neighbors)

    # Sort neighboring atoms based on distance
    distance_sort, index_sort = torch.sort(distance_sort, dim=1)
    # Select the max_num_neighbors_threshold neighbors that are closest
    distance_sort = distance_sort[:, :max_num_neighbors_threshold]
    index_sort = index_sort[:, :max_num_neighbors_threshold]

    # Offset index_sort so that it indexes into index
    index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand(
        -1, max_num_neighbors_threshold
    )
    # Remove "unused pairs" with infinite distances
    mask_finite = torch.isfinite(distance_sort)
    index_sort = torch.masked_select(index_sort, mask_finite)

    # At this point index_sort contains the index into index of the
    # closest max_num_neighbors_threshold neighbors per atom
    # Create a mask to remove all pairs not in index_sort
    mask_num_neighbors = torch.zeros(len(index), device=device, dtype=bool)
    mask_num_neighbors.index_fill_(0, index_sort, True)

    return mask_num_neighbors, num_neighbors_image


def get_pruned_edge_idx(edge_index, num_atoms=None, max_neigh=1e9):
    assert num_atoms is not None

    # removes neighbors > max_neigh
    # assumes neighbors are sorted in increasing distance
    _nonmax_idx = []
    for i in range(num_atoms):
        idx_i = torch.arange(len(edge_index[1]))[(edge_index[1] == i)][
            :max_neigh
        ]
        _nonmax_idx.append(idx_i)
    _nonmax_idx = torch.cat(_nonmax_idx)

    return _nonmax_idx


def merge_dicts(dict1: dict, dict2: dict):
    """Recursively merge two dictionaries.
    Values in dict2 override values in dict1. If dict1 and dict2 contain a dictionary as a
    value, this will call itself recursively to merge these dictionaries.
    This does not modify the input dictionaries (creates an internal copy).
    Additionally returns a list of detected duplicates.
    Adapted from https://github.com/TUM-DAML/seml/blob/master/seml/utils.py
    Parameters
    ----------
    dict1: dict
        First dict.
    dict2: dict
        Second dict. Values in dict2 will override values from dict1 in case they share the same key.
    Returns
    -------
    return_dict: dict
        Merged dictionaries.
    """
    if not isinstance(dict1, dict):
        raise ValueError(f"Expecting dict1 to be dict, found {type(dict1)}.")
    if not isinstance(dict2, dict):
        raise ValueError(f"Expecting dict2 to be dict, found {type(dict2)}.")

    return_dict = copy.deepcopy(dict1)
    duplicates = []

    for k, v in dict2.items():
        if k not in dict1:
            return_dict[k] = v
        else:
            if isinstance(v, dict) and isinstance(dict1[k], dict):
                return_dict[k], duplicates_k = merge_dicts(dict1[k], dict2[k])
                duplicates += [f"{k}.{dup}" for dup in duplicates_k]
            else:
                return_dict[k] = dict2[k]
                duplicates.append(k)

    return return_dict, duplicates


class SeverityLevelBetween(logging.Filter):
    def __init__(self, min_level, max_level):
        super().__init__()
        self.min_level = min_level
        self.max_level = max_level

    def filter(self, record):
        return self.min_level <= record.levelno < self.max_level


def setup_logging():
    root = logging.getLogger()

    # Perform setup only if logging has not been configured
    if not root.hasHandlers():
        root.setLevel(logging.INFO)

        log_formatter = logging.Formatter(
            "%(asctime)s (%(levelname)s): %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )

        # Send INFO to stdout
        handler_out = logging.StreamHandler(sys.stdout)
        handler_out.addFilter(
            SeverityLevelBetween(logging.INFO, logging.WARNING)
        )
        handler_out.setFormatter(log_formatter)
        root.addHandler(handler_out)

        # Send WARNING (and higher) to stderr
        handler_err = logging.StreamHandler(sys.stderr)
        handler_err.setLevel(logging.WARNING)
        handler_err.setFormatter(log_formatter)
        root.addHandler(handler_err)


def compute_neighbors(data, edge_index):
    # Get number of neighbors
    # segment_coo assumes sorted index
    ones = edge_index[1].new_ones(1).expand_as(edge_index[1])
    num_neighbors = segment_coo(
        ones, edge_index[1], dim_size=data["n"].sum()
    )

    # Get number of neighbors per image
    image_indptr = torch.zeros(
        data["n"].shape[0] + 1, device=data["R"].device, dtype=torch.long
    )
    image_indptr[1:] = torch.cumsum(data["n"], dim=0)
    neighbors = segment_csr(num_neighbors, image_indptr)
    return neighbors


def check_traj_files(batch, traj_dir):
    if traj_dir is None:
        return False
    traj_dir = Path(traj_dir)
    traj_files = [traj_dir / f"{id}.traj" for id in batch[0].sid.tolist()]
    return all(fl.exists() for fl in traj_files)


@contextmanager
def new_trainer_context(*, config: Dict[str, Any], args: Namespace):
    from ocpmodels.common import distutils, gp_utils
    from ocpmodels.common.registry import registry

    if TYPE_CHECKING:
        from ocpmodels.tasks.task import BaseTask
        from ocpmodels.trainers import BaseTrainer

    @dataclass
    class _TrainingContext:
        config: Dict[str, Any]
        task: "BaseTask"
        trainer: "BaseTrainer"

    setup_logging()
    original_config = config
    config = copy.deepcopy(original_config)

    if args.distributed:
        distutils.setup(config)
        if config["gp_gpus"] is not None:
            gp_utils.setup_gp(config)
    try:
        setup_imports(config)
        trainer_cls = registry.get_trainer_class(
            config.get("trainer", "energy")
        )
        assert trainer_cls is not None, "Trainer not found"
        trainer = trainer_cls(
            task=config["task"],
            model=config["model"],
            dataset=config["dataset"],
            optimizer=config["optim"],
            identifier=config["identifier"],
            timestamp_id=config.get("timestamp_id", None),
            run_dir=config.get("run_dir", "./"),
            is_debug=config.get("is_debug", False),
            print_every=config.get("print_every", 10),
            seed=config.get("seed", 0),
            logger=config.get("logger", "tensorboard"),
            local_rank=config["local_rank"],
            amp=config.get("amp", False),
            cpu=config.get("cpu", False),
            slurm=config.get("slurm", {}),
            noddp=config.get("noddp", False),
        )

        task_cls = registry.get_task_class(config["mode"])
        assert task_cls is not None, "Task not found"
        task = task_cls(config)
        start_time = time.time()
        ctx = _TrainingContext(
            config=original_config, task=task, trainer=trainer
        )
        yield ctx
        distutils.synchronize()
        if distutils.is_master():
            logging.info(f"Total time taken: {time.time() - start_time}")
    finally:
        if args.distributed:
            distutils.cleanup()


def _resolve_scale_factor_submodule(model: nn.Module, name: str):
    from ocpmodels.modules.scaling.scale_factor import ScaleFactor

    try:
        scale = model.get_submodule(name)
        if not isinstance(scale, ScaleFactor):
            return None
        return scale
    except AttributeError:
        return None


def _report_incompat_keys(
    model: nn.Module,
    keys: "_IncompatibleKeys",
    strict: bool = False,
):
    # filter out the missing scale factor keys for the new scaling factor module
    missing_keys: List[str] = []
    for full_key_name in keys.missing_keys:
        parent_module_name, _ = full_key_name.rsplit(".", 1)
        scale_factor = _resolve_scale_factor_submodule(
            model, parent_module_name
        )
        if scale_factor is not None:
            continue
        missing_keys.append(full_key_name)

    # filter out unexpected scale factor keys that remain from the old scaling modules
    unexpected_keys: List[str] = []
    for full_key_name in keys.unexpected_keys:
        parent_module_name, _ = full_key_name.rsplit(".", 1)
        scale_factor = _resolve_scale_factor_submodule(
            model, parent_module_name
        )
        if scale_factor is not None:
            continue
        unexpected_keys.append(full_key_name)

    error_msgs = []
    if len(unexpected_keys) > 0:
        error_msgs.insert(
            0,
            "Unexpected key(s) in state_dict: {}. ".format(
                ", ".join('"{}"'.format(k) for k in unexpected_keys)
            ),
        )
    if len(missing_keys) > 0:
        error_msgs.insert(
            0,
            "Missing key(s) in state_dict: {}. ".format(
                ", ".join('"{}"'.format(k) for k in missing_keys)
            ),
        )

    if len(error_msgs) > 0:
        error_msg = "Error(s) in loading state_dict for {}:\n\t{}".format(
            model.__class__.__name__, "\n\t".join(error_msgs)
        )
        if strict:
            raise RuntimeError(error_msg)
        else:
            logging.warning(error_msg)

    return missing_keys, unexpected_keys


def load_state_dict(
    module: nn.Module,
    state_dict: Mapping[str, torch.Tensor],
    strict: bool = True,
):
    incompat_keys = module.load_state_dict(state_dict, strict=False)  # type: ignore
    return _report_incompat_keys(module, incompat_keys, strict=strict)
experiments/nns/cgcnn/base_util.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import os
import sys
import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing, global_mean_pool, radius_graph
from torch_geometric.nn.models.schnet import GaussianSmearing

proj_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
if proj_root not in sys.path: sys.path.insert(1, proj_root)
from nns.cgcnn.base import BaseModel
from data.cgcnn.khot_embeddings import KHOT_EMBEDDINGS 
from data.cgcnn.qmof_khot_embeddings import QMOF_KHOT_EMBEDDINGS


class CGCNN(BaseModel):
    r"""Implementation of the Crystal Graph CNN model from the
    `"Crystal Graph Convolutional Neural Networks for an Accurate
    and Interpretable Prediction of Material Properties"
    <https://arxiv.org/abs/1710.10324>`_ paper.
    Args:
        num_atoms (int): Number of atoms.
        bond_feat_dim (int): Dimension of bond features.
        num_targets (int): Number of targets to predict.
        use_pbc (bool, optional): If set to :obj:`True`, account for periodic boundary conditions.
            (default: :obj:`True`)
        regress_forces (bool, optional): If set to :obj:`True`, predict forces by differentiating
            energy with respect to positions.
            (default: :obj:`True`)
        atom_embedding_size (int, optional): Size of atom embeddings.
            (default: :obj:`64`)
        num_graph_conv_layers (int, optional): Number of graph convolutional layers.
            (default: :obj:`6`)
        fc_feat_size (int, optional): Size of fully connected layers.
            (default: :obj:`128`)
        num_fc_layers (int, optional): Number of fully connected layers.
            (default: :obj:`4`)
        otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly.
            (default: :obj:`False`)
        cutoff (float, optional): Cutoff distance for interatomic interactions.
            (default: :obj:`10.0`)
        num_gaussians (int, optional): Number of Gaussians used for smearing.
            (default: :obj:`50.0`)
    """

    def __init__(
        self,
        num_targets=1,
        bond_feat_dim=100,
        regress_forces=True,
        atom_embedding_size=512,
        num_graph_conv_layers=3,
        fc_feat_size=128,
        num_fc_layers=3,
        cutoff=6.0,
        num_gaussians=100,
        embeddings="khot",
    ):
        super(CGCNN, self).__init__(num_targets = num_targets)
        self.regress_forces = regress_forces
        self.cutoff = cutoff
        self.max_neighbors = 50
        # Get CGCNN atom embeddings
        if embeddings == "khot":
            embeddings = KHOT_EMBEDDINGS
        elif embeddings == "qmof":
            embeddings = QMOF_KHOT_EMBEDDINGS
        else:
            raise ValueError(
                'embedding mnust be either "khot" for original CGCNN K-hot elemental embeddings or "qmof" for QMOF K-hot elemental embeddings'
            )
        self.embedding = torch.zeros(100, len(embeddings[1]))
        for i in range(100):
            self.embedding[i] = torch.tensor(embeddings[i + 1])
        self.embedding_fc = nn.Linear(len(embeddings[1]), atom_embedding_size)

        self.convs = nn.ModuleList(
            [
                CGCNNConv(
                    node_dim=atom_embedding_size,
                    edge_dim=bond_feat_dim,
                    cutoff=cutoff,
                )
                for _ in range(num_graph_conv_layers)
            ]
        )

        self.conv_to_fc = nn.Sequential(
            nn.Linear(atom_embedding_size, fc_feat_size), nn.Softplus()
        )

        if num_fc_layers > 1:
            layers = []
            for _ in range(num_fc_layers - 1):
                layers.append(nn.Linear(fc_feat_size, fc_feat_size))
                layers.append(nn.Softplus())
            self.fcs = nn.Sequential(*layers)
        self.fc_out = nn.Linear(fc_feat_size, self.num_targets)

        self.cutoff = cutoff
        self.distance_expansion = GaussianSmearing(0.0, cutoff, num_gaussians)

    def _forward(self, data):
        # Get node features
        if self.embedding.device != data["z"].device:
            self.embedding = self.embedding.to(data["z"].device)
        data["x"] = self.embedding[data["z"].long() - 1]

        edge_index, distances = self.generate_graph(data)

        data["edge_index"] = edge_index
        data["edge_attr"] = self.distance_expansion(distances)
        # Forward pass through the network
        mol_feats = self._convolve(data)
        mol_feats = self.conv_to_fc(mol_feats)
        if hasattr(self, "fcs"):
            mol_feats = self.fcs(mol_feats)

        energy = self.fc_out(mol_feats)
        return energy

    def forward(self, data):
        if self.regress_forces:
            data["R"].requires_grad_(True)
        energy = self._forward(data)

        if self.regress_forces:
            forces = -1 * (
                torch.autograd.grad(
                    energy,
                    data["R"],
                    grad_outputs=torch.ones_like(energy),
                    create_graph=True,
                )[0]
            )
            return {"E":energy, "F":forces}
        else:
            return {"E":energy}

    def _convolve(self, data):
        """
        Returns the output of the convolution layers before they are passed
        into the dense layers.
        """
        node_feats = self.embedding_fc(data["x"])
        for f in self.convs:
            node_feats = f(node_feats, data["edge_index"], data["edge_attr"])
        mol_feats = global_mean_pool(node_feats, data["batch"])
        return mol_feats
    
    def generate_graph(self, data):
        edge_index = radius_graph(
            data["R"],
            r=self.cutoff,
            batch=data["batch"],
            max_num_neighbors=self.max_neighbors,
        )

        j, i = edge_index
        distance_vec = data["R"][j] - data["R"][i]
        edge_dist = distance_vec.norm(dim=-1)

        return edge_index, edge_dist


class CGCNNConv(MessagePassing):
    """Implements the message passing layer from
    `"Crystal Graph Convolutional Neural Networks for an
    Accurate and Interpretable Prediction of Material Properties"
    <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.145301>`.
    """

    def __init__(self, node_dim, edge_dim, cutoff=6.0, **kwargs):
        super(CGCNNConv, self).__init__(aggr="add")
        self.node_feat_size = node_dim
        self.edge_feat_size = edge_dim
        self.cutoff = cutoff

        self.lin1 = nn.Linear(
            2 * self.node_feat_size + self.edge_feat_size,
            2 * self.node_feat_size,
        )
        self.bn1 = nn.BatchNorm1d(2 * self.node_feat_size)
        self.ln1 = nn.LayerNorm(self.node_feat_size)

        self.reset_parameters()

    def reset_parameters(self):
        torch.nn.init.xavier_uniform_(self.lin1.weight)

        self.lin1.bias.data.fill_(0)

        self.bn1.reset_parameters()
        self.ln1.reset_parameters()

    def forward(self, x, edge_index, edge_attr):
        """
        Arguments:
            x has shape [num_nodes, node_feat_size]
            edge_index has shape [2, num_edges]
            edge_attr is [num_edges, edge_feat_size]
        """
        out = self.propagate(
            edge_index, x=x, edge_attr=edge_attr, size=(x.size(0), x.size(0))
        )
        out = nn.Softplus()(self.ln1(out) + x)
        return out

    def message(self, x_i, x_j, edge_attr):
        """
        Arguments:
            x_i has shape [num_edges, node_feat_size]
            x_j has shape [num_edges, node_feat_size]
            edge_attr has shape [num_edges, edge_feat_size]
        Returns:
            tensor of shape [num_edges, node_feat_size]
        """
        z = self.lin1(torch.cat([x_i, x_j, edge_attr], dim=1))
        z = self.bn1(z)
        z1, z2 = z.chunk(2, dim=1)
        z1 = nn.Sigmoid()(z1)
        z2 = nn.Softplus()(z2)
        return z1 * z2


class CGCNNConv(MessagePassing):
    """Implements the message passing layer from
    `"Crystal Graph Convolutional Neural Networks for an
    Accurate and Interpretable Prediction of Material Properties"
    <https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.145301>`.
    """

    def __init__(self, node_dim, edge_dim, cutoff=6.0, **kwargs):
        super(CGCNNConv, self).__init__(aggr="add")
        self.node_feat_size = node_dim
        self.edge_feat_size = edge_dim
        self.cutoff = cutoff

        self.lin1 = nn.Linear(
            2 * self.node_feat_size + self.edge_feat_size,
            2 * self.node_feat_size,
        )
        self.bn1 = nn.BatchNorm1d(2 * self.node_feat_size)
        self.ln1 = nn.LayerNorm(self.node_feat_size)

        self.reset_parameters()

    def reset_parameters(self):
        torch.nn.init.xavier_uniform_(self.lin1.weight)

        self.lin1.bias.data.fill_(0)

        self.bn1.reset_parameters()
        self.ln1.reset_parameters()

    def forward(self, x, edge_index, edge_attr):
        """
        Arguments:
            x has shape [num_nodes, node_feat_size]
            edge_index has shape [2, num_edges]
            edge_attr is [num_edges, edge_feat_size]
        """
        out = self.propagate(
            edge_index, x=x, edge_attr=edge_attr, size=(x.size(0), x.size(0))
        )
        out = nn.Softplus()(self.ln1(out) + x)
        return out

    def message(self, x_i, x_j, edge_attr):
        """
        Arguments:
            x_i has shape [num_edges, node_feat_size]
            x_j has shape [num_edges, node_feat_size]
            edge_attr has shape [num_edges, edge_feat_size]
        Returns:
            tensor of shape [num_edges, node_feat_size]
        """
        z = self.lin1(torch.cat([x_i, x_j, edge_attr], dim=1))
        z = self.bn1(z)
        z1, z2 = z.chunk(2, dim=1)
        z1 = nn.Sigmoid()(z1)
        z2 = nn.Softplus()(z2)
        return z1 * z2
experiments/nns/cgcnn/cgcnn.py
lbai-lab/push
import torch.nn as nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        # fully connected layer, output 10 classes
        self.out = nn.Linear(32 * 7 * 7, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)

        # flatten the output of conv2 to (batch_size, 32 * 7 * 7)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output 

experiments/nns/cnn/cnn.py
lbai-lab/push
import torch

# after squeeze:
# prediction of force: [nA*batch_size, 3]
# label of force: [batch_size, nA, 3]


def EnergyLoss(pred, label):
    p, l = pred["E"].squeeze(), label["E"].squeeze()
    mae = torch.nn.L1Loss()
    return mae(p, l)


def AtomForceLoss(pred, label):
    p, l = pred["F"].squeeze(), label["F"].squeeze()
    p = p.reshape((-1, 3))
    l = l.reshape((-1, 3))
    mae = torch.nn.L1Loss()
    f = mae(p, l)
    return f


def EnergyForceLoss(pred, label):
    E = EnergyLoss(pred, label)
    F = AtomForceLoss(pred, label)
    return E + (30)*F

experiments/nns/losses.py
lbai-lab/push
import torch.nn as nn
import torchvision


class ModifiedResNet18(nn.Module):
    def __init__(self):
        super(ModifiedResNet18, self).__init__()
        self.resnet18 = torchvision.models.resnet18()
        self.resnet18.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.resnet18.fc = nn.Linear(in_features=self.resnet18.fc.in_features, out_features=10)

    def forward(self, x):
        return self.resnet18(x)

experiments/nns/resnet/resnet.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import os
import sys
import logging
import torch
import torch.nn as nn

proj_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
if proj_root not in sys.path: sys.path.insert(1, proj_root)
from torch_geometric.nn import radius_graph
from nns.schnet.base_util import (compute_neighbors, conditional_grad, get_pbc_distances, radius_graph_pbc,)


class BaseModel(nn.Module):
    def __init__(self, num_atoms=None, bond_feat_dim=None, num_targets=None):
        super(BaseModel, self).__init__()
        self.num_atoms = num_atoms
        self.bond_feat_dim = bond_feat_dim
        self.num_targets = num_targets

    def forward(self, data):
        raise NotImplementedError

    def generate_graph(
        self,
        data,
        cutoff=None,
        max_neighbors=None,
        use_pbc=None,
        otf_graph=None,
    ):
        cutoff = cutoff or self.cutoff
        max_neighbors = max_neighbors or self.max_neighbors
        use_pbc = use_pbc or self.use_pbc
        otf_graph = otf_graph or self.otf_graph

        if not otf_graph:
            try:
                edge_index = data.edge_index

                if use_pbc:
                    cell_offsets = data.cell_offsets
                    neighbors = data.neighbors

            except AttributeError:
                logging.warning(
                    "Turning otf_graph=True as required attributes not present in data object"
                )
                otf_graph = True

        if use_pbc:
            if otf_graph:
                edge_index, cell_offsets, neighbors = radius_graph_pbc(data, cutoff, max_neighbors)

            out = get_pbc_distances(
                data.pos,
                edge_index,
                data.cell,
                cell_offsets,
                neighbors,
                return_offsets=True,
                return_distance_vec=True,
            )

            edge_index = out["edge_index"]
            edge_dist = out["distances"]
            cell_offset_distances = out["offsets"]
            distance_vec = out["distance_vec"]
        else:
            if otf_graph:
                edge_index = radius_graph(data.pos, r=cutoff, batch=data.batch, max_num_neighbors=max_neighbors,)

            j, i = edge_index
            distance_vec = data.pos[j] - data.pos[i]

            edge_dist = distance_vec.norm(dim=-1)
            cell_offsets = torch.zeros(edge_index.shape[1], 3, device=data.pos.device)
            cell_offset_distances = torch.zeros_like(cell_offsets, device=data.pos.device)
            neighbors = compute_neighbors(data, edge_index)

        return (
            edge_index,
            edge_dist,
            distance_vec,
            cell_offsets,
            cell_offset_distances,
            neighbors,)

    #@property
    #def num_params(self):
    #    return sum(p.numel() for p in self.parameters())
experiments/nns/schnet/base.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

import ast
import collections
import copy
import importlib
import itertools
import json
import logging
import os
import sys
import time
from argparse import Namespace
from bisect import bisect
from contextlib import contextmanager
from dataclasses import dataclass
from functools import wraps
from itertools import product
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional

import numpy as np
import torch
import torch.nn as nn
import torch_geometric
import yaml
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from torch_geometric.data import Data
from torch_geometric.utils import remove_self_loops
from torch_scatter import segment_coo, segment_csr

if TYPE_CHECKING:
    from torch.nn.modules.module import _IncompatibleKeys


def pyg2_data_transform(data: Data):
    """
    if we're on the new pyg (2.0 or later) and if the Data stored is in older format
    we need to convert the data to the new format
    """
    if torch_geometric.__version__ >= "2.0" and "_store" not in data.__dict__:
        return Data(
            **{k: v for k, v in data.__dict__.items() if v is not None}
        )

    return data


def save_checkpoint(
    state, checkpoint_dir="checkpoints/", checkpoint_file="checkpoint.pt"
):
    filename = os.path.join(checkpoint_dir, checkpoint_file)
    torch.save(state, filename)
    return filename


class Complete(object):
    def __call__(self, data):
        device = data.edge_index.device

        row = torch.arange(data.num_nodes, dtype=torch.long, device=device)
        col = torch.arange(data.num_nodes, dtype=torch.long, device=device)

        row = row.view(-1, 1).repeat(1, data.num_nodes).view(-1)
        col = col.repeat(data.num_nodes)
        edge_index = torch.stack([row, col], dim=0)

        edge_attr = None
        if data.edge_attr is not None:
            idx = data.edge_index[0] * data.num_nodes + data.edge_index[1]
            size = list(data.edge_attr.size())
            size[0] = data.num_nodes * data.num_nodes
            edge_attr = data.edge_attr.new_zeros(size)
            edge_attr[idx] = data.edge_attr

        edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
        data.edge_attr = edge_attr
        data.edge_index = edge_index

        return data


def warmup_lr_lambda(current_step, optim_config):
    """Returns a learning rate multiplier.
    Till `warmup_steps`, learning rate linearly increases to `initial_lr`,
    and then gets multiplied by `lr_gamma` every time a milestone is crossed.
    """

    # keep this block for older configs that have warmup_epochs instead of warmup_steps
    # and lr_milestones are defined in epochs
    if (
        any(x < 100 for x in optim_config["lr_milestones"])
        or "warmup_epochs" in optim_config
    ):
        raise Exception(
            "ConfigError: please define lr_milestones in steps not epochs and define warmup_steps instead of warmup_epochs"
        )

    if current_step <= optim_config["warmup_steps"]:
        alpha = current_step / float(optim_config["warmup_steps"])
        return optim_config["warmup_factor"] * (1.0 - alpha) + alpha
    else:
        idx = bisect(optim_config["lr_milestones"], current_step)
        return pow(optim_config["lr_gamma"], idx)


def print_cuda_usage():
    print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 * 1024))
    print(
        "Max Memory Allocated:",
        torch.cuda.max_memory_allocated() / (1024 * 1024),
    )
    print("Memory Cached:", torch.cuda.memory_cached() / (1024 * 1024))
    print("Max Memory Cached:", torch.cuda.max_memory_cached() / (1024 * 1024))


def conditional_grad(dec):
    "Decorator to enable/disable grad depending on whether force/energy predictions are being made"
    # Adapted from https://stackoverflow.com/questions/60907323/accessing-class-property-as-decorator-argument
    def decorator(func):
        @wraps(func)
        def cls_method(self, *args, **kwargs):
            f = func
            if self.regress_forces and not getattr(self, "direct_forces", 0):
                f = dec(func)
            return f(self, *args, **kwargs)

        return cls_method

    return decorator


def plot_histogram(data, xlabel="", ylabel="", title=""):
    assert isinstance(data, list)

    # Preset
    fig = Figure(figsize=(5, 4), dpi=150)
    canvas = FigureCanvas(fig)
    ax = fig.gca()

    # Plot
    ax.hist(data, bins=20, rwidth=0.9, zorder=3)

    # Axes
    ax.grid(color="0.95", zorder=0)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    fig.tight_layout(pad=2)

    # Return numpy array
    canvas.draw()
    image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    image_from_plot = image_from_plot.reshape(
        fig.canvas.get_width_height()[::-1] + (3,)
    )

    return image_from_plot


# Override the collation method in `pytorch_geometric.data.InMemoryDataset`
def collate(data_list):
    keys = data_list[0].keys
    data = data_list[0].__class__()

    for key in keys:
        data[key] = []
    slices = {key: [0] for key in keys}

    for item, key in product(data_list, keys):
        data[key].append(item[key])
        if torch.is_tensor(item[key]):
            s = slices[key][-1] + item[key].size(
                item.__cat_dim__(key, item[key])
            )
        elif isinstance(item[key], int) or isinstance(item[key], float):
            s = slices[key][-1] + 1
        else:
            raise ValueError("Unsupported attribute type")
        slices[key].append(s)

    if hasattr(data_list[0], "__num_nodes__"):
        data.__num_nodes__ = []
        for item in data_list:
            data.__num_nodes__.append(item.num_nodes)

    for key in keys:
        if torch.is_tensor(data_list[0][key]):
            data[key] = torch.cat(
                data[key], dim=data.__cat_dim__(key, data_list[0][key])
            )
        else:
            data[key] = torch.tensor(data[key])
        slices[key] = torch.tensor(slices[key], dtype=torch.long)

    return data, slices


def add_edge_distance_to_graph(
    batch,
    device="cpu",
    dmin=0.0,
    dmax=6.0,
    num_gaussians=50,
):
    # Make sure x has positions.
    if not all(batch.pos[0][:] == batch.x[0][-3:]):
        batch.x = torch.cat([batch.x, batch.pos.float()], dim=1)
    # First set computations to be tracked for positions.
    batch.x = batch.x.requires_grad_(True)
    # Then compute Euclidean distance between edge endpoints.
    pdist = torch.nn.PairwiseDistance(p=2.0)
    distances = pdist(
        batch.x[batch.edge_index[0]][:, -3:],
        batch.x[batch.edge_index[1]][:, -3:],
    )
    # Expand it using a gaussian basis filter.
    gdf_filter = torch.linspace(dmin, dmax, num_gaussians)
    var = gdf_filter[1] - gdf_filter[0]
    gdf_filter, var = gdf_filter.to(device), var.to(device)
    gdf_distances = torch.exp(
        -((distances.view(-1, 1) - gdf_filter) ** 2) / var**2
    )
    # Reassign edge attributes.
    batch.edge_weight = distances
    batch.edge_attr = gdf_distances.float()
    return batch


def _import_local_file(path: Path, *, project_root: Path):
    """
    Imports a Python file as a module
    :param path: The path to the file to import
    :type path: Path
    :param project_root: The root directory of the project (i.e., the "ocp" folder)
    :type project_root: Path
    """

    path = path.resolve()
    project_root = project_root.resolve()

    module_name = ".".join(
        path.absolute()
        .relative_to(project_root.absolute())
        .with_suffix("")
        .parts
    )
    logging.debug(f"Resolved module name of {path} to {module_name}")
    importlib.import_module(module_name)


def setup_experimental_imports(project_root: Path):
    experimental_folder = (project_root / "experimental").resolve()
    if not experimental_folder.exists() or not experimental_folder.is_dir():
        return

    experimental_files = [
        f.resolve().absolute() for f in experimental_folder.rglob("*.py")
    ]
    # Ignore certain directories within experimental
    ignore_file = experimental_folder / ".ignore"
    if ignore_file.exists():
        with open(ignore_file, "r") as f:
            for line in f.read().splitlines():
                for ignored_file in (experimental_folder / line).rglob("*.py"):
                    experimental_files.remove(
                        ignored_file.resolve().absolute()
                    )

    for f in experimental_files:
        _import_local_file(f, project_root=project_root)


def _get_project_root():
    """
    Gets the root folder of the project (the "ocp" folder)
    :return: The absolute path to the project root.
    """
    from ocpmodels.common.registry import registry

    # Automatically load all of the modules, so that
    # they register with registry
    root_folder = registry.get("ocpmodels_root", no_warning=True)

    if root_folder is not None:
        assert isinstance(root_folder, str), "ocpmodels_root must be a string"
        root_folder = Path(root_folder).resolve().absolute()
        assert root_folder.exists(), f"{root_folder} does not exist"
        assert root_folder.is_dir(), f"{root_folder} is not a directory"
    else:
        root_folder = Path(__file__).resolve().absolute().parent.parent

    # root_folder is the "ocpmodes" folder, so we need to go up one more level
    return root_folder.parent


# Copied from https://github.com/facebookresearch/mmf/blob/master/mmf/utils/env.py#L89.
def setup_imports(config: Optional[dict] = None):
    from ocpmodels.common.registry import registry

    skip_experimental_imports = (config or {}).get(
        "skip_experimental_imports", None
    )

    # First, check if imports are already setup
    has_already_setup = registry.get("imports_setup", no_warning=True)
    if has_already_setup:
        return

    try:
        project_root = _get_project_root()
        logging.info(f"Project root: {project_root}")
        importlib.import_module("ocpmodels.common.logger")

        import_keys = ["trainers", "datasets", "models", "tasks"]
        for key in import_keys:
            for f in (project_root / "ocpmodels" / key).rglob("*.py"):
                _import_local_file(f, project_root=project_root)

        if not skip_experimental_imports:
            setup_experimental_imports(project_root)
    finally:
        registry.register("imports_setup", True)


def dict_set_recursively(dictionary, key_sequence, val):
    top_key = key_sequence.pop(0)
    if len(key_sequence) == 0:
        dictionary[top_key] = val
    else:
        if top_key not in dictionary:
            dictionary[top_key] = {}
        dict_set_recursively(dictionary[top_key], key_sequence, val)


def parse_value(value):
    """
    Parse string as Python literal if possible and fallback to string.
    """
    try:
        return ast.literal_eval(value)
    except (ValueError, SyntaxError):
        # Use as string if nothing else worked
        return value


def create_dict_from_args(args: list, sep: str = "."):
    """
    Create a (nested) dictionary from console arguments.
    Keys in different dictionary levels are separated by sep.
    """
    return_dict = {}
    for arg in args:
        arg = arg.strip("--")
        keys_concat, val = arg.split("=")
        val = parse_value(val)
        key_sequence = keys_concat.split(sep)
        dict_set_recursively(return_dict, key_sequence, val)
    return return_dict


def load_config(path: str, previous_includes: list = []):
    path = Path(path)
    if path in previous_includes:
        raise ValueError(
            f"Cyclic config include detected. {path} included in sequence {previous_includes}."
        )
    previous_includes = previous_includes + [path]

    direct_config = yaml.safe_load(open(path, "r"))

    # Load config from included files.
    if "includes" in direct_config:
        includes = direct_config.pop("includes")
    else:
        includes = []
    if not isinstance(includes, list):
        raise AttributeError(
            "Includes must be a list, '{}' provided".format(type(includes))
        )

    config = {}
    duplicates_warning = []
    duplicates_error = []

    for include in includes:
        include_config, inc_dup_warning, inc_dup_error = load_config(
            include, previous_includes
        )
        duplicates_warning += inc_dup_warning
        duplicates_error += inc_dup_error

        # Duplicates between includes causes an error
        config, merge_dup_error = merge_dicts(config, include_config)
        duplicates_error += merge_dup_error

    # Duplicates between included and main file causes warnings
    config, merge_dup_warning = merge_dicts(config, direct_config)
    duplicates_warning += merge_dup_warning

    return config, duplicates_warning, duplicates_error


def build_config(args, args_override):
    config, duplicates_warning, duplicates_error = load_config(args.config_yml)
    if len(duplicates_warning) > 0:
        logging.warning(
            f"Overwritten config parameters from included configs "
            f"(non-included parameters take precedence): {duplicates_warning}"
        )
    if len(duplicates_error) > 0:
        raise ValueError(
            f"Conflicting (duplicate) parameters in simultaneously "
            f"included configs: {duplicates_error}"
        )

    # Check for overridden parameters.
    if args_override != []:
        overrides = create_dict_from_args(args_override)
        config, _ = merge_dicts(config, overrides)

    # Some other flags.
    config["mode"] = args.mode
    config["identifier"] = args.identifier
    config["timestamp_id"] = args.timestamp_id
    config["seed"] = args.seed
    config["is_debug"] = args.debug
    config["run_dir"] = args.run_dir
    config["print_every"] = args.print_every
    config["amp"] = args.amp
    config["checkpoint"] = args.checkpoint
    config["cpu"] = args.cpu
    # Submit
    config["submit"] = args.submit
    config["summit"] = args.summit
    # Distributed
    config["local_rank"] = args.local_rank
    config["distributed_port"] = args.distributed_port
    config["world_size"] = args.num_nodes * args.num_gpus
    config["distributed_backend"] = args.distributed_backend
    config["noddp"] = args.no_ddp
    config["gp_gpus"] = args.gp_gpus

    return config


def create_grid(base_config, sweep_file):
    def _flatten_sweeps(sweeps, root_key="", sep="."):
        flat_sweeps = []
        for key, value in sweeps.items():
            new_key = root_key + sep + key if root_key else key
            if isinstance(value, collections.MutableMapping):
                flat_sweeps.extend(_flatten_sweeps(value, new_key).items())
            else:
                flat_sweeps.append((new_key, value))
        return collections.OrderedDict(flat_sweeps)

    def _update_config(config, keys, override_vals, sep="."):
        for key, value in zip(keys, override_vals):
            key_path = key.split(sep)
            child_config = config
            for name in key_path[:-1]:
                child_config = child_config[name]
            child_config[key_path[-1]] = value
        return config

    sweeps = yaml.safe_load(open(sweep_file, "r"))
    flat_sweeps = _flatten_sweeps(sweeps)
    keys = list(flat_sweeps.keys())
    values = list(itertools.product(*flat_sweeps.values()))

    configs = []
    for i, override_vals in enumerate(values):
        config = copy.deepcopy(base_config)
        config = _update_config(config, keys, override_vals)
        config["identifier"] = config["identifier"] + f"_run{i}"
        configs.append(config)
    return configs


def save_experiment_log(args, jobs, configs):
    log_file = args.logdir / "exp" / time.strftime("%Y-%m-%d-%I-%M-%S%p.log")
    log_file.parent.mkdir(exist_ok=True, parents=True)
    with open(log_file, "w") as f:
        for job, config in zip(jobs, configs):
            print(
                json.dumps(
                    {
                        "config": config,
                        "slurm_id": job.job_id,
                        "timestamp": time.strftime("%I:%M:%S%p %Z %b %d, %Y"),
                    }
                ),
                file=f,
            )
    return log_file


def get_pbc_distances(
    pos,
    edge_index,
    cell,
    cell_offsets,
    neighbors,
    return_offsets=False,
    return_distance_vec=False,
):
    row, col = edge_index

    distance_vectors = pos[row] - pos[col]

    # correct for pbc
    neighbors = neighbors.to(cell.device)
    cell = torch.repeat_interleave(cell, neighbors, dim=0)
    offsets = cell_offsets.float().view(-1, 1, 3).bmm(cell.float()).view(-1, 3)
    distance_vectors += offsets

    # compute distances
    distances = distance_vectors.norm(dim=-1)

    # redundancy: remove zero distances
    nonzero_idx = torch.arange(len(distances))[distances != 0]
    edge_index = edge_index[:, nonzero_idx]
    distances = distances[nonzero_idx]

    out = {
        "edge_index": edge_index,
        "distances": distances,
    }

    if return_distance_vec:
        out["distance_vec"] = distance_vectors[nonzero_idx]

    if return_offsets:
        out["offsets"] = offsets[nonzero_idx]

    return out


def radius_graph_pbc(
    data, radius, max_num_neighbors_threshold, pbc=[True, True, False]
):
    device = data.pos.device
    batch_size = len(data.natoms)

    # position of the atoms
    atom_pos = data.pos

    # Before computing the pairwise distances between atoms, first create a list of atom indices to compare for the entire batch
    num_atoms_per_image = data.natoms
    num_atoms_per_image_sqr = (num_atoms_per_image**2).long()

    # index offset between images
    index_offset = (
        torch.cumsum(num_atoms_per_image, dim=0) - num_atoms_per_image
    )

    index_offset_expand = torch.repeat_interleave(
        index_offset, num_atoms_per_image_sqr
    )
    num_atoms_per_image_expand = torch.repeat_interleave(
        num_atoms_per_image, num_atoms_per_image_sqr
    )

    # Compute a tensor containing sequences of numbers that range from 0 to num_atoms_per_image_sqr for each image
    # that is used to compute indices for the pairs of atoms. This is a very convoluted way to implement
    # the following (but 10x faster since it removes the for loop)
    # for batch_idx in range(batch_size):
    #    batch_count = torch.cat([batch_count, torch.arange(num_atoms_per_image_sqr[batch_idx], device=device)], dim=0)
    num_atom_pairs = torch.sum(num_atoms_per_image_sqr)
    index_sqr_offset = (
        torch.cumsum(num_atoms_per_image_sqr, dim=0) - num_atoms_per_image_sqr
    )
    index_sqr_offset = torch.repeat_interleave(
        index_sqr_offset, num_atoms_per_image_sqr
    )
    atom_count_sqr = (
        torch.arange(num_atom_pairs, device=device) - index_sqr_offset
    )

    # Compute the indices for the pairs of atoms (using division and mod)
    # If the systems get too large this apporach could run into numerical precision issues
    index1 = (
        torch.div(
            atom_count_sqr, num_atoms_per_image_expand, rounding_mode="floor"
        )
    ) + index_offset_expand
    index2 = (
        atom_count_sqr % num_atoms_per_image_expand
    ) + index_offset_expand
    # Get the positions for each atom
    pos1 = torch.index_select(atom_pos, 0, index1)
    pos2 = torch.index_select(atom_pos, 0, index2)

    # Calculate required number of unit cells in each direction.
    # Smallest distance between planes separated by a1 is
    # 1 / ||(a2 x a3) / V||_2, since a2 x a3 is the area of the plane.
    # Note that the unit cell volume V = a1 * (a2 x a3) and that
    # (a2 x a3) / V is also the reciprocal primitive vector
    # (crystallographer's definition).

    cross_a2a3 = torch.cross(data.cell[:, 1], data.cell[:, 2], dim=-1)
    cell_vol = torch.sum(data.cell[:, 0] * cross_a2a3, dim=-1, keepdim=True)

    if pbc[0]:
        inv_min_dist_a1 = torch.norm(cross_a2a3 / cell_vol, p=2, dim=-1)
        rep_a1 = torch.ceil(radius * inv_min_dist_a1)
    else:
        rep_a1 = data.cell.new_zeros(1)

    if pbc[1]:
        cross_a3a1 = torch.cross(data.cell[:, 2], data.cell[:, 0], dim=-1)
        inv_min_dist_a2 = torch.norm(cross_a3a1 / cell_vol, p=2, dim=-1)
        rep_a2 = torch.ceil(radius * inv_min_dist_a2)
    else:
        rep_a2 = data.cell.new_zeros(1)

    if pbc[2]:
        cross_a1a2 = torch.cross(data.cell[:, 0], data.cell[:, 1], dim=-1)
        inv_min_dist_a3 = torch.norm(cross_a1a2 / cell_vol, p=2, dim=-1)
        rep_a3 = torch.ceil(radius * inv_min_dist_a3)
    else:
        rep_a3 = data.cell.new_zeros(1)

    # Take the max over all images for uniformity. This is essentially padding.
    # Note that this can significantly increase the number of computed distances
    # if the required repetitions are very different between images
    # (which they usually are). Changing this to sparse (scatter) operations
    # might be worth the effort if this function becomes a bottleneck.
    max_rep = [rep_a1.max(), rep_a2.max(), rep_a3.max()]

    # Tensor of unit cells
    cells_per_dim = [
        torch.arange(-rep, rep + 1, device=device, dtype=torch.float)
        for rep in max_rep
    ]
    unit_cell = torch.cartesian_prod(*cells_per_dim)
    num_cells = len(unit_cell)
    unit_cell_per_atom = unit_cell.view(1, num_cells, 3).repeat(
        len(index2), 1, 1
    )
    unit_cell = torch.transpose(unit_cell, 0, 1)
    unit_cell_batch = unit_cell.view(1, 3, num_cells).expand(
        batch_size, -1, -1
    )

    # Compute the x, y, z positional offsets for each cell in each image
    data_cell = torch.transpose(data.cell, 1, 2)
    pbc_offsets = torch.bmm(data_cell, unit_cell_batch)
    pbc_offsets_per_atom = torch.repeat_interleave(
        pbc_offsets, num_atoms_per_image_sqr, dim=0
    )

    # Expand the positions and indices for the 9 cells
    pos1 = pos1.view(-1, 3, 1).expand(-1, -1, num_cells)
    pos2 = pos2.view(-1, 3, 1).expand(-1, -1, num_cells)
    index1 = index1.view(-1, 1).repeat(1, num_cells).view(-1)
    index2 = index2.view(-1, 1).repeat(1, num_cells).view(-1)
    # Add the PBC offsets for the second atom
    pos2 = pos2 + pbc_offsets_per_atom

    # Compute the squared distance between atoms
    atom_distance_sqr = torch.sum((pos1 - pos2) ** 2, dim=1)
    atom_distance_sqr = atom_distance_sqr.view(-1)

    # Remove pairs that are too far apart
    mask_within_radius = torch.le(atom_distance_sqr, radius * radius)
    # Remove pairs with the same atoms (distance = 0.0)
    mask_not_same = torch.gt(atom_distance_sqr, 0.0001)
    mask = torch.logical_and(mask_within_radius, mask_not_same)
    index1 = torch.masked_select(index1, mask)
    index2 = torch.masked_select(index2, mask)
    unit_cell = torch.masked_select(
        unit_cell_per_atom.view(-1, 3), mask.view(-1, 1).expand(-1, 3)
    )
    unit_cell = unit_cell.view(-1, 3)
    atom_distance_sqr = torch.masked_select(atom_distance_sqr, mask)

    mask_num_neighbors, num_neighbors_image = get_max_neighbors_mask(
        natoms=data.natoms,
        index=index1,
        atom_distance=atom_distance_sqr,
        max_num_neighbors_threshold=max_num_neighbors_threshold,
    )

    if not torch.all(mask_num_neighbors):
        # Mask out the atoms to ensure each atom has at most max_num_neighbors_threshold neighbors
        index1 = torch.masked_select(index1, mask_num_neighbors)
        index2 = torch.masked_select(index2, mask_num_neighbors)
        unit_cell = torch.masked_select(
            unit_cell.view(-1, 3), mask_num_neighbors.view(-1, 1).expand(-1, 3)
        )
        unit_cell = unit_cell.view(-1, 3)

    edge_index = torch.stack((index2, index1))

    return edge_index, unit_cell, num_neighbors_image


def get_max_neighbors_mask(
    natoms, index, atom_distance, max_num_neighbors_threshold
):
    """
    Give a mask that filters out edges so that each atom has at most
    `max_num_neighbors_threshold` neighbors.
    Assumes that `index` is sorted.
    """
    device = natoms.device
    num_atoms = natoms.sum()

    # Get number of neighbors
    # segment_coo assumes sorted index
    ones = index.new_ones(1).expand_as(index)
    num_neighbors = segment_coo(ones, index, dim_size=num_atoms)
    max_num_neighbors = num_neighbors.max()
    num_neighbors_thresholded = num_neighbors.clamp(
        max=max_num_neighbors_threshold
    )

    # Get number of (thresholded) neighbors per image
    image_indptr = torch.zeros(
        natoms.shape[0] + 1, device=device, dtype=torch.long
    )
    image_indptr[1:] = torch.cumsum(natoms, dim=0)
    num_neighbors_image = segment_csr(num_neighbors_thresholded, image_indptr)

    # If max_num_neighbors is below the threshold, return early
    if (
        max_num_neighbors <= max_num_neighbors_threshold
        or max_num_neighbors_threshold <= 0
    ):
        mask_num_neighbors = torch.tensor(
            [True], dtype=bool, device=device
        ).expand_as(index)
        return mask_num_neighbors, num_neighbors_image

    # Create a tensor of size [num_atoms, max_num_neighbors] to sort the distances of the neighbors.
    # Fill with infinity so we can easily remove unused distances later.
    distance_sort = torch.full(
        [num_atoms * max_num_neighbors], np.inf, device=device
    )

    # Create an index map to map distances from atom_distance to distance_sort
    # index_sort_map assumes index to be sorted
    index_neighbor_offset = torch.cumsum(num_neighbors, dim=0) - num_neighbors
    index_neighbor_offset_expand = torch.repeat_interleave(
        index_neighbor_offset, num_neighbors
    )
    index_sort_map = (
        index * max_num_neighbors
        + torch.arange(len(index), device=device)
        - index_neighbor_offset_expand
    )
    distance_sort.index_copy_(0, index_sort_map, atom_distance)
    distance_sort = distance_sort.view(num_atoms, max_num_neighbors)

    # Sort neighboring atoms based on distance
    distance_sort, index_sort = torch.sort(distance_sort, dim=1)
    # Select the max_num_neighbors_threshold neighbors that are closest
    distance_sort = distance_sort[:, :max_num_neighbors_threshold]
    index_sort = index_sort[:, :max_num_neighbors_threshold]

    # Offset index_sort so that it indexes into index
    index_sort = index_sort + index_neighbor_offset.view(-1, 1).expand(
        -1, max_num_neighbors_threshold
    )
    # Remove "unused pairs" with infinite distances
    mask_finite = torch.isfinite(distance_sort)
    index_sort = torch.masked_select(index_sort, mask_finite)

    # At this point index_sort contains the index into index of the
    # closest max_num_neighbors_threshold neighbors per atom
    # Create a mask to remove all pairs not in index_sort
    mask_num_neighbors = torch.zeros(len(index), device=device, dtype=bool)
    mask_num_neighbors.index_fill_(0, index_sort, True)

    return mask_num_neighbors, num_neighbors_image


def get_pruned_edge_idx(edge_index, num_atoms=None, max_neigh=1e9):
    assert num_atoms is not None

    # removes neighbors > max_neigh
    # assumes neighbors are sorted in increasing distance
    _nonmax_idx = []
    for i in range(num_atoms):
        idx_i = torch.arange(len(edge_index[1]))[(edge_index[1] == i)][
            :max_neigh
        ]
        _nonmax_idx.append(idx_i)
    _nonmax_idx = torch.cat(_nonmax_idx)

    return _nonmax_idx


def merge_dicts(dict1: dict, dict2: dict):
    """Recursively merge two dictionaries.
    Values in dict2 override values in dict1. If dict1 and dict2 contain a dictionary as a
    value, this will call itself recursively to merge these dictionaries.
    This does not modify the input dictionaries (creates an internal copy).
    Additionally returns a list of detected duplicates.
    Adapted from https://github.com/TUM-DAML/seml/blob/master/seml/utils.py
    Parameters
    ----------
    dict1: dict
        First dict.
    dict2: dict
        Second dict. Values in dict2 will override values from dict1 in case they share the same key.
    Returns
    -------
    return_dict: dict
        Merged dictionaries.
    """
    if not isinstance(dict1, dict):
        raise ValueError(f"Expecting dict1 to be dict, found {type(dict1)}.")
    if not isinstance(dict2, dict):
        raise ValueError(f"Expecting dict2 to be dict, found {type(dict2)}.")

    return_dict = copy.deepcopy(dict1)
    duplicates = []

    for k, v in dict2.items():
        if k not in dict1:
            return_dict[k] = v
        else:
            if isinstance(v, dict) and isinstance(dict1[k], dict):
                return_dict[k], duplicates_k = merge_dicts(dict1[k], dict2[k])
                duplicates += [f"{k}.{dup}" for dup in duplicates_k]
            else:
                return_dict[k] = dict2[k]
                duplicates.append(k)

    return return_dict, duplicates


class SeverityLevelBetween(logging.Filter):
    def __init__(self, min_level, max_level):
        super().__init__()
        self.min_level = min_level
        self.max_level = max_level

    def filter(self, record):
        return self.min_level <= record.levelno < self.max_level


def setup_logging():
    root = logging.getLogger()

    # Perform setup only if logging has not been configured
    if not root.hasHandlers():
        root.setLevel(logging.INFO)

        log_formatter = logging.Formatter(
            "%(asctime)s (%(levelname)s): %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )

        # Send INFO to stdout
        handler_out = logging.StreamHandler(sys.stdout)
        handler_out.addFilter(
            SeverityLevelBetween(logging.INFO, logging.WARNING)
        )
        handler_out.setFormatter(log_formatter)
        root.addHandler(handler_out)

        # Send WARNING (and higher) to stderr
        handler_err = logging.StreamHandler(sys.stderr)
        handler_err.setLevel(logging.WARNING)
        handler_err.setFormatter(log_formatter)
        root.addHandler(handler_err)


def compute_neighbors(data, edge_index):
    # Get number of neighbors
    # segment_coo assumes sorted index
    ones = edge_index[1].new_ones(1).expand_as(edge_index[1])
    num_neighbors = segment_coo(
        ones, edge_index[1], dim_size=data["n"].sum()
    )

    # Get number of neighbors per image
    image_indptr = torch.zeros(
        data["n"].shape[0] + 1, device=data["R"].device, dtype=torch.long
    )
    image_indptr[1:] = torch.cumsum(data["n"], dim=0)
    neighbors = segment_csr(num_neighbors, image_indptr)
    return neighbors


def check_traj_files(batch, traj_dir):
    if traj_dir is None:
        return False
    traj_dir = Path(traj_dir)
    traj_files = [traj_dir / f"{id}.traj" for id in batch[0].sid.tolist()]
    return all(fl.exists() for fl in traj_files)


@contextmanager
def new_trainer_context(*, config: Dict[str, Any], args: Namespace):
    from ocpmodels.common import distutils, gp_utils
    from ocpmodels.common.registry import registry

    if TYPE_CHECKING:
        from ocpmodels.tasks.task import BaseTask
        from ocpmodels.trainers import BaseTrainer

    @dataclass
    class _TrainingContext:
        config: Dict[str, Any]
        task: "BaseTask"
        trainer: "BaseTrainer"

    setup_logging()
    original_config = config
    config = copy.deepcopy(original_config)

    if args.distributed:
        distutils.setup(config)
        if config["gp_gpus"] is not None:
            gp_utils.setup_gp(config)
    try:
        setup_imports(config)
        trainer_cls = registry.get_trainer_class(
            config.get("trainer", "energy")
        )
        assert trainer_cls is not None, "Trainer not found"
        trainer = trainer_cls(
            task=config["task"],
            model=config["model"],
            dataset=config["dataset"],
            optimizer=config["optim"],
            identifier=config["identifier"],
            timestamp_id=config.get("timestamp_id", None),
            run_dir=config.get("run_dir", "./"),
            is_debug=config.get("is_debug", False),
            print_every=config.get("print_every", 10),
            seed=config.get("seed", 0),
            logger=config.get("logger", "tensorboard"),
            local_rank=config["local_rank"],
            amp=config.get("amp", False),
            cpu=config.get("cpu", False),
            slurm=config.get("slurm", {}),
            noddp=config.get("noddp", False),
        )

        task_cls = registry.get_task_class(config["mode"])
        assert task_cls is not None, "Task not found"
        task = task_cls(config)
        start_time = time.time()
        ctx = _TrainingContext(
            config=original_config, task=task, trainer=trainer
        )
        yield ctx
        distutils.synchronize()
        if distutils.is_master():
            logging.info(f"Total time taken: {time.time() - start_time}")
    finally:
        if args.distributed:
            distutils.cleanup()


def _resolve_scale_factor_submodule(model: nn.Module, name: str):
    from ocpmodels.modules.scaling.scale_factor import ScaleFactor

    try:
        scale = model.get_submodule(name)
        if not isinstance(scale, ScaleFactor):
            return None
        return scale
    except AttributeError:
        return None


def _report_incompat_keys(
    model: nn.Module,
    keys: "_IncompatibleKeys",
    strict: bool = False,
):
    # filter out the missing scale factor keys for the new scaling factor module
    missing_keys: List[str] = []
    for full_key_name in keys.missing_keys:
        parent_module_name, _ = full_key_name.rsplit(".", 1)
        scale_factor = _resolve_scale_factor_submodule(
            model, parent_module_name
        )
        if scale_factor is not None:
            continue
        missing_keys.append(full_key_name)

    # filter out unexpected scale factor keys that remain from the old scaling modules
    unexpected_keys: List[str] = []
    for full_key_name in keys.unexpected_keys:
        parent_module_name, _ = full_key_name.rsplit(".", 1)
        scale_factor = _resolve_scale_factor_submodule(
            model, parent_module_name
        )
        if scale_factor is not None:
            continue
        unexpected_keys.append(full_key_name)

    error_msgs = []
    if len(unexpected_keys) > 0:
        error_msgs.insert(
            0,
            "Unexpected key(s) in state_dict: {}. ".format(
                ", ".join('"{}"'.format(k) for k in unexpected_keys)
            ),
        )
    if len(missing_keys) > 0:
        error_msgs.insert(
            0,
            "Missing key(s) in state_dict: {}. ".format(
                ", ".join('"{}"'.format(k) for k in missing_keys)
            ),
        )

    if len(error_msgs) > 0:
        error_msg = "Error(s) in loading state_dict for {}:\n\t{}".format(
            model.__class__.__name__, "\n\t".join(error_msgs)
        )
        if strict:
            raise RuntimeError(error_msg)
        else:
            logging.warning(error_msg)

    return missing_keys, unexpected_keys


def load_state_dict(
    module: nn.Module,
    state_dict: Mapping[str, torch.Tensor],
    strict: bool = True,
):
    incompat_keys = module.load_state_dict(state_dict, strict=False)  # type: ignore
    return _report_incompat_keys(module, incompat_keys, strict=strict)
experiments/nns/schnet/base_util.py
lbai-lab/push
"""
Copyright (c) Facebook, Inc. and its affiliates.
This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""
import os
import sys
import torch

proj_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
if proj_root not in sys.path: sys.path.insert(1, proj_root)
from torch_geometric.nn import SchNet
from nns.schnet.base import BaseModel

class SchNetWrap(SchNet, BaseModel):
    r"""Wrapper around the continuous-filter convolutional neural network SchNet from the
    `"SchNet: A Continuous-filter Convolutional Neural Network for Modeling
    Quantum Interactions" <https://arxiv.org/abs/1706.08566>`_. Each layer uses interaction
    block of the form:
    .. math::
        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \odot
        h_{\mathbf{\Theta}} ( \exp(-\gamma(\mathbf{e}_{j,i} - \mathbf{\mu}))),
    Args:
        num_atoms (int): Unused argument
        bond_feat_dim (int): Unused argument
        num_targets (int): Number of targets to predict.
        use_pbc (bool, optional): If set to :obj:`True`, account for periodic boundary conditions.
            (default: :obj:`True`)
        regress_forces (bool, optional): If set to :obj:`True`, predict forces by differentiating
            energy with respect to positions.
            (default: :obj:`True`)
        otf_graph (bool, optional): If set to :obj:`True`, compute graph edges on the fly.
            (default: :obj:`False`)
        hidden_channels (int, optional): Number of hidden channels.
            (default: :obj:`128`)
        num_filters (int, optional): Number of filters to use.
            (default: :obj:`128`)
        num_interactions (int, optional): Number of interaction blocks
            (default: :obj:`6`)
        num_gaussians (int, optional): The number of gaussians :math:`\mu`.
            (default: :obj:`50`)
        cutoff (float, optional): Cutoff distance for interatomic interactions.
            (default: :obj:`10.0`)
        readout (string, optional): Whether to apply :obj:`"add"` or
            :obj:`"mean"` global aggregation. (default: :obj:`"add"`)
    """

    def __init__(
        self,
        regress_forces=True,
        hidden_channels=1024,
        num_filters=256,
        num_interactions=5,
        num_gaussians=200,
        cutoff=6.0,
        readout="add",):

        self.regress_forces = regress_forces
        self.cutoff = cutoff
        self.max_neighbors = 50
        self.freeze_layers = ['embedding.weight', 'interactions.0.mlp.0.weight', 'interactions.0.mlp.2.weight', 'interactions.0.conv.lin1.weight', 'interactions.0.conv.lin2.weight', 
                              'interactions.0.conv.lin2.bias', 'interactions.0.lin.weight', 'interactions.0.lin.bias', 'interactions.1.mlp.0.weight', 'interactions.1.mlp.2.weight', 
                              'interactions.1.conv.lin1.weight', 'interactions.1.conv.lin2.weight', 'interactions.1.conv.lin2.bias', 'interactions.1.lin.weight', 
                              'interactions.1.lin.bias', 'interactions.2.mlp.0.weight', 'interactions.2.mlp.2.weight', 'interactions.2.conv.lin1.weight', 
                              'interactions.2.conv.lin2.weight', 'interactions.2.conv.lin2.bias', 'interactions.2.lin.weight', 'interactions.2.lin.bias', 
                              'interactions.3.mlp.0.weight', 'interactions.3.mlp.2.weight', 'interactions.3.conv.lin1.weight', 'interactions.3.conv.lin2.weight', 
                              'interactions.3.conv.lin2.bias', 'interactions.3.lin.weight', 'interactions.3.lin.bias', 'interactions.4.conv.lin1.weight', 
                              'interactions.4.conv.lin2.weight', 'interactions.4.conv.lin2.bias', 'interactions.4.lin.weight', 'interactions.4.lin.bias', 
                              'lin1.weight', 'lin1.bias', 'lin2.bias']

        super(SchNetWrap, self).__init__(
            hidden_channels=hidden_channels,
            num_filters=num_filters,
            num_interactions=num_interactions,
            num_gaussians=num_gaussians,
            cutoff=cutoff,
            readout=readout,)

    def forward(self, data):
        z = data["z"]
        pos = data["R"]
        batch = data["batch"]

        if self.regress_forces:
            pos.requires_grad_(True)
        energy = super(SchNetWrap, self).forward(z, pos, batch=batch)

        if self.regress_forces:
            forces = -1*(torch.autograd.grad(energy, pos, grad_outputs=torch.ones_like(energy), create_graph=True,)[0])
            return {"E":energy, "F":forces}
        else:
            return {"E":energy}

experiments/nns/schnet/schnet.py
lbai-lab/push
# -*- coding: utf-8 -*-
"""
       <NAME OF THE PROGRAM THIS FILE BELONGS TO>

  File:     utils.py
  Authors:  Timothy Praditia (timothy.praditia@iws.uni-stuttgart.de)
            Raphael Leiteritz (raphael.leiteritz@ipvs.uni-stuttgart.de)
            Makoto Takamoto (makoto.takamoto@neclab.eu)
            Francesco Alesiani (makoto.takamoto@neclab.eu)

NEC Laboratories Europe GmbH, Copyright (c) <year>, All rights reserved.

       THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY.

       PROPRIETARY INFORMATION ---

SOFTWARE LICENSE AGREEMENT

ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY

BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS
LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR
DOWNLOAD THE SOFTWARE.

This is a license agreement ("Agreement") between your academic institution
or non-profit organization or self (called "Licensee" or "You" in this
Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this
Agreement).  All rights not specifically granted to you in this Agreement
are reserved for Licensor.

RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive
ownership of any copy of the Software (as defined below) licensed under this
Agreement and hereby grants to Licensee a personal, non-exclusive,
non-transferable license to use the Software for noncommercial research
purposes, without the right to sublicense, pursuant to the terms and
conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF
LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this
Agreement, the term "Software" means (i) the actual copy of all or any
portion of code for program routines made accessible to Licensee by Licensor
pursuant to this Agreement, inclusive of backups, updates, and/or merged
copies permitted hereunder or subsequently supplied by Licensor,  including
all or any file structures, programming instructions, user interfaces and
screen formats and sequences as well as any and all documentation and
instructions related to it, and (ii) all or any derivatives and/or
modifications created or made by You to any of the items specified in (i).

CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is
proprietary to Licensor, and as such, Licensee agrees to receive all such
materials and to use the Software only in accordance with the terms of this
Agreement.  Licensee agrees to use reasonable effort to protect the Software
from unauthorized use, reproduction, distribution, or publication. All
publication materials mentioning features or use of this software must
explicitly include an acknowledgement the software was developed by NEC
Laboratories Europe GmbH.

COPYRIGHT: The Software is owned by Licensor.

PERMITTED USES:  The Software may be used for your own noncommercial
internal research purposes. You understand and agree that Licensor is not
obligated to implement any suggestions and/or feedback you might provide
regarding the Software, but to the extent Licensor does so, you are not
entitled to any compensation related thereto.

DERIVATIVES: You may create derivatives of or make modifications to the
Software, however, You agree that all and any such derivatives and
modifications will be owned by Licensor and become a part of the Software
licensed to You under this Agreement.  You may only use such derivatives and
modifications for your own noncommercial internal research purposes, and you
may not otherwise use, distribute or copy such derivatives and modifications
in violation of this Agreement.

BACKUPS:  If Licensee is an organization, it may make that number of copies
of the Software necessary for internal noncommercial use at a single site
within its organization provided that all information appearing in or on the
original labels, including the copyright and trademark notices are copied
onto the labels of the copies.

USES NOT PERMITTED:  You may not distribute, copy or use the Software except
as explicitly permitted herein. Licensee has not been granted any trademark
license as part of this Agreement.  Neither the name of NEC Laboratories
Europe GmbH nor the names of its contributors may be used to endorse or
promote products derived from this Software without specific prior written
permission.

You may not sell, rent, lease, sublicense, lend, time-share or transfer, in
whole or in part, or provide third parties access to prior or present
versions (or any parts thereof) of the Software.

ASSIGNMENT: You may not assign this Agreement or your rights hereunder
without the prior written consent of Licensor. Any attempted assignment
without such consent shall be null and void.

TERM: The term of the license granted by this Agreement is from Licensee's
acceptance of this Agreement by downloading the Software or by using the
Software until terminated as provided below.

The Agreement automatically terminates without notice if you fail to comply
with any provision of this Agreement.  Licensee may terminate this Agreement
by ceasing using the Software.  Upon any termination of this Agreement,
Licensee will delete any and all copies of the Software. You agree that all
provisions which operate to protect the proprietary rights of Licensor shall
remain in force should breach occur and that the obligation of
confidentiality described in this Agreement is binding in perpetuity and, as
such, survives the term of the Agreement.

FEE: Provided Licensee abides completely by the terms and conditions of this
Agreement, there is no fee due to Licensor for Licensee's use of the
Software in accordance with this Agreement.

DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY
OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR
FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT.  LICENSEE
BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND
RELATED MATERIALS.

SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is
provided as part of this Agreement.

EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent
permitted under applicable law, Licensor shall not be liable for direct,
indirect, special, incidental, or consequential damages or lost profits
related to Licensee's use of and/or inability to use the Software, even if
Licensor is advised of the possibility of such damage.

EXPORT REGULATION: Licensee agrees to comply with any and all applicable
export control laws, regulations, and/or other laws related to embargoes and
sanction programs administered by law.

SEVERABILITY: If any provision(s) of this Agreement shall be held to be
invalid, illegal, or unenforceable by a court or other tribunal of competent
jurisdiction, the validity, legality and enforceability of the remaining
provisions shall not in any way be affected or impaired thereby.

NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right
or remedy under this Agreement shall be construed as a waiver of any future
or other exercise of such right or remedy by Licensor.

GOVERNING LAW: This Agreement shall be construed and enforced in accordance
with the laws of Germany without reference to conflict of laws principles.
You consent to the personal jurisdiction of the courts of this country and
waive their rights to venue outside of Germany.

ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and
entire agreement between Licensee and Licensor as to the matter set forth
herein and supersedes any previous agreements, understandings, and
arrangements between the parties relating hereto.

       THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY.
"""

import torch
from torch.utils.data import Dataset
import os
import h5py
import numpy as np
import math as mt


class UNetDatasetSingle(Dataset):
    def __init__(self, filename,
                 initial_step=10,
                 saved_folder='../data/',
                 reduced_resolution=1,
                 reduced_resolution_t=1,
                 reduced_batch=1,
                 if_test=False,
                 test_ratio=0.1,
                 num_samples_max = -1):
        """
        
        :param filename: filename that contains the dataset
        :type filename: STR
        :param filenum: array containing indices of filename included in the dataset
        :type filenum: ARRAY

        """
        
        # Define path to files
        root_path = os.path.abspath(saved_folder + filename)
        print("ROOT_PATH", root_path)
        assert filename[-2:] != 'h5', 'HDF5 data is assumed!!'
        
        with h5py.File(root_path, 'r') as f:
            keys = list(f.keys())
            keys.sort()
            if 'tensor' not in keys:
                _data = np.array(f['density'], dtype=np.float32)  # batch, time, x,...
                idx_cfd = _data.shape
                if len(idx_cfd)==3:  # 1D
                    self.data = np.zeros([idx_cfd[0]//reduced_batch,
                                          idx_cfd[2]//reduced_resolution,
                                          mt.ceil(idx_cfd[1]/reduced_resolution_t),
                                          3],
                                         dtype=np.float32)
                    #density
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :], (0, 2, 1))
                    self.data[...,0] = _data   # batch, x, t, ch
                    # pressure
                    _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :], (0, 2, 1))
                    self.data[...,1] = _data   # batch, x, t, ch
                    # Vx
                    _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :], (0, 2, 1))
                    self.data[...,2] = _data   # batch, x, t, ch

                if len(idx_cfd)==4:  # 2D
                    self.data = np.zeros([idx_cfd[0]//reduced_batch,
                                          idx_cfd[2]//reduced_resolution,
                                          idx_cfd[3]//reduced_resolution,
                                          mt.ceil(idx_cfd[1]/reduced_resolution_t),
                                          4],
                                         dtype=np.float32)
                    # density
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 1))
                    self.data[...,0] = _data   # batch, x, t, ch
                    # pressure
                    _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 1))
                    self.data[...,1] = _data   # batch, x, t, ch
                    # Vx
                    _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 1))
                    self.data[...,2] = _data   # batch, x, t, ch
                    # Vy
                    _data = np.array(f['Vy'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 1))
                    self.data[...,3] = _data   # batch, x, t, ch
                    
                if len(idx_cfd)==5:  # 3D
                    self.data = np.zeros([idx_cfd[0]//reduced_batch,
                                          idx_cfd[2]//reduced_resolution,
                                          idx_cfd[3]//reduced_resolution,
                                          idx_cfd[4]//reduced_resolution,
                                          mt.ceil(idx_cfd[1]/reduced_resolution_t),
                                          5],
                                         dtype=np.float32)
                    # density
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 4, 1))
                    self.data[...,0] = _data   # batch, x, t, ch
                    # pressure
                    _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 4, 1))
                    self.data[...,1] = _data   # batch, x, t, ch
                    # Vx
                    _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 4, 1))
                    self.data[...,2] = _data   # batch, x, t, ch
                    # Vy
                    _data = np.array(f['Vy'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 4, 1))
                    self.data[...,3] = _data   # batch, x, t, ch
                    # Vz
                    _data = np.array(f['Vz'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data, (0, 2, 3, 4, 1))
                    self.data[...,4] = _data   # batch, x, t, ch

            else:  # scalar equations
                ## data dim = [t, x1, ..., xd, v]
                _data = np.array(f['tensor'], dtype=np.float32)  # batch, time, x,...
                if len(_data.shape) == 3:  # 1D
                    _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :], (0, 2, 1))
                    self.data = _data[:, :, :, None]  # batch, x, t, ch

                if len(_data.shape) == 4:  # 2D Darcy flow
                    # u: label
                    _data = _data[::reduced_batch,:,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1))
                    #if _data.shape[-1]==1:  # if nt==1
                    #    _data = np.tile(_data, (1, 1, 1, 2))
                    self.data = _data
                    # nu: input
                    _data = np.array(f['nu'], dtype=np.float32)  # batch, time, x,...
                    _data = _data[::reduced_batch, None,::reduced_resolution,::reduced_resolution]
                    ## convert to [x1, ..., xd, t, v]
                    _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1))
                    self.data = np.concatenate([_data, self.data], axis=-1)
                    self.data = self.data[:, :, :, :, None]  # batch, x, y, t, ch

        if num_samples_max>0:
            num_samples_max  = min(num_samples_max,self.data.shape[0])
        else:
            num_samples_max = self.data.shape[0]

        test_idx = int(num_samples_max * test_ratio)
        if if_test:
            self.data = self.data[:test_idx]
        else:
            self.data = self.data[test_idx:num_samples_max]

        # Time steps used as initial conditions
        self.initial_step = initial_step
        
        self.data = torch.tensor(self.data)


    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        
        return self.data[idx,...,:self.initial_step,:], self.data[idx]


    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        
        return self.data[idx,...,:self.initial_step,:], self.data[idx]

# class UNetDatasetSingle(Dataset):
#     def __init__(self, filename,device,
#                  initial_step=10,
#                  saved_folder='../data/',
#                  reduced_resolution=1,
#                  reduced_resolution_t=1,
#                  reduced_batch=1,
#                  if_test=False,
#                  test_ratio=0.1,
#                  num_samples_max = -1):
#         """
        
#         :param filename: filename that contains the dataset
#         :type filename: STR
#         :param filenum: array containing indices of filename included in the dataset
#         :type filenum: ARRAY

#         """
#         self.device=device
#         # Define path to files
#         root_path = os.path.abspath(saved_folder + filename)
#         assert filename[-2:] != 'h5', 'HDF5 data is assumed!!'
        
#         with h5py.File(root_path, 'r') as f:
#             # keys = list(f.keys())
#             # keys.sort()

#             # if 'tensor' not in keys:
#             #     _data = np.array(f['density'], dtype=np.float32)  # batch, time, x,...
#             #     idx_cfd = _data.shape
#             #     if len(idx_cfd)==3:  # 1D
#             #         self.data = np.zeros([idx_cfd[0]//reduced_batch,
#             #                               idx_cfd[2]//reduced_resolution,
#             #                               mt.ceil(idx_cfd[1]/reduced_resolution_t),
#             #                               3],
#             #                              dtype=np.float32)
#             #         #density
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data[:, :, :], (0, 2, 1))
#             #         self.data[...,0] = _data   # batch, x, t, ch
#             #         # pressure
#             #         _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data[:, :, :], (0, 2, 1))
#             #         self.data[...,1] = _data   # batch, x, t, ch
#             #         # Vx
#             #         _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data[:, :, :], (0, 2, 1))
#             #         self.data[...,2] = _data   # batch, x, t, ch

#             #     if len(idx_cfd)==4:  # 2D
#             #         self.data = np.zeros([idx_cfd[0]//reduced_batch,
#             #                               idx_cfd[2]//reduced_resolution,
#             #                               idx_cfd[3]//reduced_resolution,
#             #                               mt.ceil(idx_cfd[1]/reduced_resolution_t),
#             #                               4],
#             #                              dtype=np.float32)
#             #         # density
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 1))
#             #         self.data[...,0] = _data   # batch, x, t, ch
#             #         # pressure
#             #         _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 1))
#             #         self.data[...,1] = _data   # batch, x, t, ch
#             #         # Vx
#             #         _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 1))
#             #         self.data[...,2] = _data   # batch, x, t, ch
#             #         # Vy
#             #         _data = np.array(f['Vy'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 1))
#             #         self.data[...,3] = _data   # batch, x, t, ch
                    
#             #     if len(idx_cfd)==5:  # 3D
#             #         self.data = np.zeros([idx_cfd[0]//reduced_batch,
#             #                               idx_cfd[2]//reduced_resolution,
#             #                               idx_cfd[3]//reduced_resolution,
#             #                               idx_cfd[4]//reduced_resolution,
#             #                               mt.ceil(idx_cfd[1]/reduced_resolution_t),
#             #                               5],
#             #                              dtype=np.float32)
#             #         # density
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 4, 1))
#             #         self.data[...,0] = _data   # batch, x, t, ch
#             #         # pressure
#             #         _data = np.array(f['pressure'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 4, 1))
#             #         self.data[...,1] = _data   # batch, x, t, ch
#             #         # Vx
#             #         _data = np.array(f['Vx'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 4, 1))
#             #         self.data[...,2] = _data   # batch, x, t, ch
#             #         # Vy
#             #         _data = np.array(f['Vy'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 4, 1))
#             #         self.data[...,3] = _data   # batch, x, t, ch
#             #         # Vz
#             #         _data = np.array(f['Vz'], dtype=np.float32)  # batch, time, x,...
#             #         _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution,::reduced_resolution]
#             #         ## convert to [x1, ..., xd, t, v]
#             #         _data = np.transpose(_data, (0, 2, 3, 4, 1))
#             #         self.data[...,4] = _data   # batch, x, t, ch

#             #else:  # scalar equations
#                 ## data dim = [t, x1, ..., xd, v]
#             _data = np.array(f['tensor'], dtype=np.float32)  # batch, time, x,...
#             if len(_data.shape) == 3:  # 1D
#                 _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution]
#                 ## convert to [x1, ..., xd, t, v]
#                 _data = np.transpose(_data[:, :, :], (0, 2, 1))
#                 self.data = _data[:, :, :, None]  # batch, x, t, ch

#                 # if len(_data.shape) == 4:  # 2D Darcy flow
#                 #     # u: label
#                 #     _data = _data[::reduced_batch,:,::reduced_resolution,::reduced_resolution]
#                 #     ## convert to [x1, ..., xd, t, v]
#                 #     _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1))
#                 #     #if _data.shape[-1]==1:  # if nt==1
#                 #     #    _data = np.tile(_data, (1, 1, 1, 2))
#                 #     self.data = _data
#                 #     # nu: input
#                 #     _data = np.array(f['nu'], dtype=np.float32)  # batch, time, x,...
#                 #     _data = _data[::reduced_batch, None,::reduced_resolution,::reduced_resolution]
#                 #     ## convert to [x1, ..., xd, t, v]
#                 #     _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1))
#                 #     self.data = np.concatenate([_data, self.data], axis=-1)
#                 #     self.data = self.data[:, :, :, :, None]  # batch, x, y, t, ch

#         if num_samples_max>0:
#             num_samples_max  = min(num_samples_max,self.data.shape[0])
#         else:
#             num_samples_max = self.data.shape[0]

#         test_idx = int(num_samples_max * test_ratio)
#         if if_test:
#             self.data = self.data[:test_idx]
#         else:
#             self.data = self.data[test_idx:num_samples_max]

#         # Time steps used as initial conditions
#         self.initial_step = initial_step
        
#         self.data = torch.tensor(self.data).to(self.device)
#     def __len__(self):
#         return len(self.data)
    
#     def __getitem__(self, idx):
#         return self.data[idx,...,:self.initial_step,:], self.data[idx]
    

class UNetDatasetMult(Dataset):
    def __init__(self, filename,
                 initial_step=10,
                 saved_folder='../data/',
                 reduced_resolution=1,
                 reduced_resolution_t=1,
                 reduced_batch=1,
                 if_test=False, test_ratio=0.1
                 ):
        """
        
        :param filename: filename that contains the dataset
        :type filename: STR
        :param filenum: array containing indices of filename included in the dataset
        :type filenum: ARRAY
        :param initial_step: time steps taken as initial condition, defaults to 10
        :type initial_step: INT, optional

        """
        
        # Define path to files
        self.file_path = os.path.abspath(saved_folder + filename + ".h5")
        
        # Extract list of seeds
        with h5py.File(self.file_path, 'r') as h5_file:
            data_list = sorted(h5_file.keys())

        test_idx = int(len(data_list) * (1-test_ratio))
        if if_test:
            self.data_list = np.array(data_list[test_idx:])
        else:
            self.data_list = np.array(data_list[:test_idx])
        
        # Time steps used as initial conditions
        self.initial_step = initial_step

    def __len__(self):
        return len(self.data_list)
    
    def __getitem__(self, idx):
        
        # Open file and read data
        with h5py.File(self.file_path, 'r') as h5_file:
            seed_group = h5_file[self.data_list[idx]]
        
            # data dim = [t, x1, ..., xd, v]
            data = np.array(seed_group["data"], dtype='f')
            data = torch.tensor(data, dtype=torch.float)
            
            # convert to [x1, ..., xd, t, v]
            permute_idx = list(range(1,len(data.shape)-1))
            permute_idx.extend(list([0, -1]))
            data = data.permute(permute_idx)
            
        return data[...,:self.initial_step,:], data
experiments/nns/unet/UnetDataSet.py
lbai-lab/push
#!/usr/bin/env python3

"""
U-Net. Implementation taken and modified from
https://github.com/mateuszbuda/brain-segmentation-pytorch

MIT License

Copyright (c) 2019 mateuszbuda

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"""
from collections import OrderedDict

import torch
import torch.nn as nn


class UNet1d(nn.Module):

    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet1d, self).__init__()

        features = init_features
        self.encoder1 = UNet1d._block(in_channels, features, name="enc1")
        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.encoder2 = UNet1d._block(features, features * 2, name="enc2")
        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.encoder3 = UNet1d._block(features * 2, features * 4, name="enc3")
        self.pool3 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.encoder4 = UNet1d._block(features * 4, features * 8, name="enc4")
        self.pool4 = nn.MaxPool1d(kernel_size=2, stride=2)

        self.bottleneck = UNet1d._block(features * 8, features * 16, name="bottleneck")

        self.upconv4 = nn.ConvTranspose1d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = UNet1d._block((features * 8) * 2, features * 8, name="dec4")
        self.upconv3 = nn.ConvTranspose1d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = UNet1d._block((features * 4) * 2, features * 4, name="dec3")
        self.upconv2 = nn.ConvTranspose1d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = UNet1d._block((features * 2) * 2, features * 2, name="dec2")
        self.upconv1 = nn.ConvTranspose1d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = UNet1d._block(features * 2, features, name="dec1")

        self.conv = nn.Conv1d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return self.conv(dec1)

    @staticmethod
    def _block(in_channels, features, name):
        return nn.Sequential(
            OrderedDict(
                [
                    (
                        name + "conv1",
                        nn.Conv1d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm1", nn.BatchNorm1d(num_features=features)),
                    (name + "tanh1", nn.Tanh()),
                    (
                        name + "conv2",
                        nn.Conv1d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm2", nn.BatchNorm1d(num_features=features)),
                    (name + "tanh2", nn.Tanh()),
                ]
            )
        )


class UNet2d(nn.Module):

    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet2d, self).__init__()

        features = init_features
        self.encoder1 = UNet2d._block(in_channels, features, name="enc1")
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = UNet2d._block(features, features * 2, name="enc2")
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = UNet2d._block(features * 2, features * 4, name="enc3")
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = UNet2d._block(features * 4, features * 8, name="enc4")
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.bottleneck = UNet2d._block(features * 8, features * 16, name="bottleneck")

        self.upconv4 = nn.ConvTranspose2d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = UNet2d._block((features * 8) * 2, features * 8, name="dec4")
        self.upconv3 = nn.ConvTranspose2d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = UNet2d._block((features * 4) * 2, features * 4, name="dec3")
        self.upconv2 = nn.ConvTranspose2d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = UNet2d._block((features * 2) * 2, features * 2, name="dec2")
        self.upconv1 = nn.ConvTranspose2d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = UNet2d._block(features * 2, features, name="dec1")

        self.conv = nn.Conv2d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return self.conv(dec1)

    @staticmethod
    def _block(in_channels, features, name):
        return nn.Sequential(
            OrderedDict(
                [
                    (
                        name + "conv1",
                        nn.Conv2d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm1", nn.BatchNorm2d(num_features=features)),
                    (name + "tanh1", nn.Tanh()),
                    (
                        name + "conv2",
                        nn.Conv2d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm2", nn.BatchNorm2d(num_features=features)),
                    (name + "tanh2", nn.Tanh()),
                ]
            )
        )
    

class UNet3d(nn.Module):

    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet3d, self).__init__()

        features = init_features
        self.encoder1 = UNet3d._block(in_channels, features, name="enc1")
        self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2)
        self.encoder2 = UNet3d._block(features, features * 2, name="enc2")
        self.pool2 = nn.MaxPool3d(kernel_size=2, stride=2)
        self.encoder3 = UNet3d._block(features * 2, features * 4, name="enc3")
        self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2)
        self.encoder4 = UNet3d._block(features * 4, features * 8, name="enc4")
        self.pool4 = nn.MaxPool3d(kernel_size=2, stride=2)

        self.bottleneck = UNet3d._block(features * 8, features * 16, name="bottleneck")

        self.upconv4 = nn.ConvTranspose3d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = UNet3d._block((features * 8) * 2, features * 8, name="dec4")
        self.upconv3 = nn.ConvTranspose3d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = UNet3d._block((features * 4) * 2, features * 4, name="dec3")
        self.upconv2 = nn.ConvTranspose3d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = UNet3d._block((features * 2) * 2, features * 2, name="dec2")
        self.upconv1 = nn.ConvTranspose3d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = UNet3d._block(features * 2, features, name="dec1")

        self.conv = nn.Conv3d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return self.conv(dec1)

    @staticmethod
    def _block(in_channels, features, name):
        return nn.Sequential(
            OrderedDict(
                [
                    (
                        name + "conv1",
                        nn.Conv3d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm1", nn.BatchNorm3d(num_features=features)),
                    (name + "tanh1", nn.Tanh()),
                    (
                        name + "conv2",
                        nn.Conv3d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding=1,
                            bias=False,
                        ),
                    ),
                    (name + "norm2", nn.BatchNorm3d(num_features=features)),
                    (name + "tanh2", nn.Tanh()),
                ]
            )
        )
experiments/nns/unet/unet.py
lbai-lab/push

import torch.nn as nn
from nns.unet.unet import UNet1d


class UNet1dWrap(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.nn = UNet1d(in_channels, out_channels)

    def forward(self, x):
        return self.nn.forward(x).permute([0, 2, 1])


def unet_loss_fn(x, y):
    _batch = y.size(0)
    loss = nn.MSELoss()(x.reshape(_batch, -1), y.reshape(_batch, -1))
    return loss


def unet_loss_fn2(x, y):
    _batch = y.size(0)
    loss = nn.MSELoss()(x.reshape(_batch, -1), y.reshape(_batch, -1))
    return loss

experiments/nns/unet/unetWrap.py
lbai-lab/push
import torchvision
import torch.nn as nn


class Modifiedvit_b_16(nn.Module):
    def __init__(self):
        super(Modifiedvit_b_16, self).__init__()
        # self.vit_b_16 = torchvision.models.VisionTransformer(image_size=28, patch_size=14, num_classes=10, num_heads=4, num_layers=20, mlp_dim=1024, hidden_dim=128)
        self.vit_b_16 = torchvision.models.VisionTransformer(image_size=28, patch_size=14, num_classes=10, num_heads=8, num_layers=16, mlp_dim=1280, hidden_dim=320)

    def forward(self, x):
        return self.vit_b_16(x)


class Modifiedvit_b_16_adjust(nn.Module):
    def __init__(self, num_heads, num_layers, mlp_dim, hidden_dim):
        super(Modifiedvit_b_16_adjust, self).__init__()
        # self.vit_b_16 = torchvision.models.VisionTransformer(image_size=28, patch_size=14, num_classes=10, num_heads=4, num_layers=20, mlp_dim=1024, hidden_dim=128)
        self.vit_b_16 = torchvision.models.VisionTransformer(image_size=28, patch_size=14, num_classes=10, num_heads=num_heads, num_layers=num_layers, mlp_dim=mlp_dim, hidden_dim=hidden_dim)

    def forward(self, x):
        return self.vit_b_16(x)

experiments/nns/vit/vit.py
lbai-lab/push
import argparse
from functools import partial
import numpy as np
from typing import *

import push.bayes.ensemble
import push.bayes.swag
import push.bayes.stein_vgd

from infer.standard import train_standard
from infer.ensemble import train_deep_ensemble
from infer.mswag import train_mswag
from infer.stein_vgd import train_svgd
from infer.push_ensemble import _ensemble_main_instrumented, mk_optim
from infer.push_mswag import _mswag_particle_instrumented, _mswag_sample_entry_instrumented, _mswag_sample_instrumented
from infer.push_stein_vgd import _svgd_leader_instrumented, _svgd_leader_instrumented_memeff

from train_util import training_params, get_model, get_model_and_args


if __name__ == "__main__":

    # Initialize a list to store memory usage for each GPU
    parser = argparse.ArgumentParser()
    
    # [Training params]
    parser.add_argument("-wb", "--wandb", action='store_true')
    parser.add_argument("-g", "--group", type=str, default='default')
    parser.add_argument("-e", "--epochs", type=int, default=5)
    parser.add_argument("-lr", "--learning_rate", type=float, default=0.001)
    parser.add_argument("-dl", "--dataset_length", type=int, default=None)
    parser.add_argument("-w", "--weight_decay", type=float, default=0)
    parser.add_argument("-ss", "--scheduler_step", type=int, default=100)
    parser.add_argument("-sg", "--scheduler_gamma", type=int, default=0.5)
    parser.add_argument("-model", "--model", type=str, default="vit",
                        choices=[
                            "schnet",
                            "cgcnn",
                            "unet",
                            "resnet",
                            "cnn",
                            "transformer",
                            "transformer2",
                        ])
    parser.add_argument("-t", "--train", type=str, default="svgd_push",
                        choices=[
                            "standard",
                            "ensemble",
                            "ensemble_push",
                            "mswag",
                            "mswag_push",
                            "svgd",
                            "svgd_push"
                        ])
    
    # [Push args]
    parser.add_argument("-n", "--num_particles", type=int, default=2)
    parser.add_argument("-cs", "--cache_size", type=int, default=None)
    parser.add_argument("-vs", "--view_size", type=int, default=None)
    parser.add_argument("-d", "--devices", type=int, default=1)
    parser.add_argument("-save", "--save", action="store_true")
    parser.add_argument("-ppred", "--posterior-pred", action="store_true")

    # [SWAG params]
    parser.add_argument("-pe", "--pretrain_epochs", type=int, default=1)
    parser.add_argument("-se", "--swag_epochs", type=int, default=1)
    parser.add_argument("-samples", "--samples", type=int, default=5)
    parser.add_argument("-scale", "--scale", type=float, default=1.0)
    parser.add_argument("-var_clamp", "--var_clamp", type=float, default=1e-30)

    # [SVGD params]
    parser.add_argument("-band", "--bandwidth", type=float, default=1)
    parser.add_argument("-mef", "--mem_eff", action="store_true")
    
    # Add training parameters
    args, train_loader, test_loader, loss_fn = training_params(parser)

    # GPU memory
    if args.cache_size is None:
        args.cache_size = int(np.ceil(args.num_particles/args.devices))
    if args.view_size is None:
        args.view_size = args.num_particles
        

    # Select
    if args.train == "standard":
        trained_model = train_standard(train_loader, args, loss_fn)

    elif args.train == "ensemble":
        trained_models = train_deep_ensemble(train_loader, args, loss_fn)

    elif args.train == "ensemble_push":
        model, model_args = get_model_and_args(args)
        ensemble_state = {
            "args": args,
        }
        push.bayes.ensemble.train_deep_ensemble(
            train_loader, loss_fn, args.epochs, model, *model_args,
            cache_size=args.cache_size, view_size=args.view_size, 
            num_ensembles=args.num_particles, num_devices=args.devices,
            mk_optim=partial(mk_optim, args.learning_rate, args.weight_decay),
            ensemble_entry=_ensemble_main_instrumented, ensemble_state=ensemble_state
        )
    
    elif args.train == "mswag":
        nns = train_mswag(train_loader, args, loss_fn)

    elif args.train == "mswag_push":
        model, model_args = get_model_and_args(args)
        if args.cache_size is None:
            args.cache_size = np.ceil(args.num_particles/args.devices)
        # args.view_size = args.num_particles
        mswag_state = { 
            "args": args,
        }
        mswag = push.bayes.swag.train_mswag(
            train_loader, loss_fn, args.pretrain_epochs, args.swag_epochs,
            args.num_particles, args.cache_size, args.view_size, 
            model, *model_args,
            num_devices=args.devices, lr=args.learning_rate,
            mswag_entry=_mswag_particle_instrumented, mswag_state=mswag_state,
            f_save=args.save, mswag_sample_entry=_mswag_sample_entry_instrumented, mswag_sample=_mswag_sample_instrumented
        )
        if args.posterior_pred:
            mswag.posterior_pred(test_loader, loss_fn, num_samples=args.samples , scale=args.scale, var_clamp=args.var_clamp)
    
    elif args.train == "svgd": 
        networks = train_svgd(train_loader, args, loss_fn) 

    elif args.train == "svgd_push":
        prior = None
        svgd_state = { 
            "args": args,
        }
        model, model_args = get_model_and_args(args)
        if args.mem_eff:
            svgd_entry = _svgd_leader_instrumented_memeff
        else:
            svgd_entry = _svgd_leader_instrumented
        push.bayes.stein_vgd.train_svgd(
            train_loader, loss_fn, args.epochs, args.num_particles, model, *model_args,
            num_devices=args.devices, cache_size=args.cache_size, view_size=args.view_size,
            lengthscale=args.bandwidth, lr=args.learning_rate, prior=prior,
            svgd_entry=svgd_entry, svgd_state=svgd_state
        )
    else:
        raise ValueError(f"Method {args.train} not supported ...")

experiments/train.py
lbai-lab/push
from datetime import datetime
import pytz
from functools import partial
import torch
import torch.nn as nn
from timeit import default_timer
import time
import wandb

# [Qchem]
from nns.schnet.schnet import SchNetWrap
from nns.cgcnn.cgcnn import CGCNN
from nns.losses import EnergyForceLoss

# [Sciml]
from nns.unet.UnetDataSet import UNetDatasetSingle
from nns.unet.unetWrap import UNet1dWrap, unet_loss_fn

# [Vision]
from nns.vit.vit import Modifiedvit_b_16, Modifiedvit_b_16_adjust
from nns.cnn.cnn import CNN
from nns.resnet.resnet import ModifiedResNet18

# [Datasets]
from torch.utils.data import DataLoader
from data.MD17.MD17Dataset import MD17SingleDataset
import torchvision.datasets as datasets
from torchvision.transforms import ToTensor
import torchvision.transforms as transforms


# =============================================================================
# Dataloader Collate
# =============================================================================

def unet_collate(batch, args):
    data, label = zip(*batch)
    
    data = torch.stack(data)
    label = torch.stack(label)
    
    x = data[..., 0, :]
    y = label[..., args.t_train - 1:args.t_train, :]
    
    x = x.permute([0, 2, 1])
    
    return x, y
    

def qchem_collate(batch):
    data = {"R":[], "z":[], "batch":[], "n":[]}
    label = {"E":[], "F":[], "covariance_inv":batch[0][4], "mean":batch[0][5]}      

    for i, b in enumerate(batch):
        data["R"].append(b[0])
        data["z"].append(b[1])
        data["n"].append(len(b[1]))
        data["batch"].append(torch.ones((b[1].size()), dtype=torch.int64)*i)
        label["E"].append(b[2])
        label["F"].append(b[3])
    data["R"] = torch.cat(data["R"])
    data["z"] = torch.cat(data["z"])
    data["n"] = torch.Tensor(data["n"])
    data["batch"] = torch.cat(data["batch"])
    label["E"] = torch.stack(label["E"])
    label["F"] = torch.stack(label["F"])
    
    data = {i: v for i, v in data.items()}
    label = {i: v for i, v in label.items()}
    
    return data, label


def vision_collate(batch):
    # Unzip the batch into data and labels lists
    data_list, label_list = zip(*batch)
    
    data = torch.stack(data_list)
    label = torch.tensor(label_list, dtype=torch.long)
        
    return data, label


# =============================================================================
# Model selection
# =============================================================================

def get_model_and_args(args):
    if args.nn_module == UNet1dWrap:
        return UNet1dWrap, [args.in_channels, args.out_channels]
    elif args.nn_module == SchNetWrap:
        return SchNetWrap, []
    elif args.nn_module == CGCNN:
        return CGCNN, []
    elif args.nn_module == CNN:
        return CNN, []
    elif args.nn_module == ModifiedResNet18:
        return ModifiedResNet18, []
    elif args.nn_module == Modifiedvit_b_16:
        return Modifiedvit_b_16, []
    elif args.nn_module == Modifiedvit_b_16_adjust:
        import numpy as np
        args.num_params = sum([np.prod(p.size()) for p in Modifiedvit_b_16_adjust(args.num_heads, args.num_layers, args.mlp_dim, args.hidden_dim).parameters()])
        print("PARAMETERS", args.num_params)
        return Modifiedvit_b_16_adjust, [args.num_heads, args.num_layers, args.mlp_dim, args.hidden_dim]
    else:
        raise ValueError(f"Model {args.nn_module} not supported ...")


def get_model(args):
    model = None
    if args.nn_module == UNet1dWrap:
        model  = UNet1dWrap(args.in_channels, args.out_channels)
    
    elif args.nn_module == SchNetWrap:
        model = SchNetWrap()
    
    elif args.nn_module == CGCNN:
        model = CGCNN()
    
    elif args.nn_module == CNN:
        model = CNN()

    elif args.nn_module == ModifiedResNet18:
        model = ModifiedResNet18()

    elif args.nn_module == Modifiedvit_b_16:
        model = Modifiedvit_b_16()

    elif args.nn_module == Modifiedvit_b_16_adjust:
        model = Modifiedvit_b_16_adjust(args.num_heads, args.num_layers, args.mlp_dim, args.hidden_dim)

    return model


# =============================================================================
# Loss function selection
# =============================================================================

def get_loss_fn(args):
    if args.model == "unet":
        return unet_loss_fn
    elif args.model == "schnet" or args.model == "cgcnn":
        return EnergyForceLoss
    elif args.model == "cnn" or  args.model == "resnet" or args.model == "transformer" or args.model == "transformer2":
        return nn.CrossEntropyLoss()
    else:
        raise ValueError(f"Model {args.model} not supported ...")


# =============================================================================
# Dataloader selection
# =============================================================================


def get_dataloaders(args):
    print("\nGenerating Data Loaders.....")
    if args.model == "unet":
        file_name = '1D_Advection_Sols_beta0.1.hdf5'
        if args.cloud_path:
            base_path = '/home/paperspace/PusH2/experiments/data/1D/Advection/Train/'
        else:
            base_path = '/usr/data1/PDEBench/pdebench/data/1D/Advection/'
        train_dataset = UNetDatasetSingle(file_name,
                                          saved_folder=base_path,
                                          reduced_resolution=args.reduced_resolution,
                                          reduced_resolution_t=args.reduced_resolution_t,
                                          reduced_batch=args.reduced_batch,
                                          initial_step=args.initial_step)

        test_dataset = UNetDatasetSingle(file_name,
                                         saved_folder=base_path,
                                         reduced_resolution=args.reduced_resolution,
                                         reduced_resolution_t=args.reduced_resolution_t,
                                         reduced_batch=args.reduced_batch,
                                         initial_step=args.initial_step,
                                         if_test=True)
        
        if args.dataset_length is not None:
            train_dataset = torch.utils.data.random_split(train_dataset, [args.dataset_length, len(train_dataset) - args.dataset_length])[0]
            test_dataset = torch.utils.data.random_split(test_dataset, [int(args.dataset_length/10), len(test_dataset) - int(args.dataset_length/10)])[0]
            print("truncated train_dataset", len(train_dataset))

        # Create a decoy dataloader to update t_train 
        temp_train_loader = DataLoader(train_dataset, batch_size=1, num_workers=0, shuffle=False)
        _, _data = next(iter(temp_train_loader))
        dimensions = len(_data.shape)
        print("Spatial Dimension", dimensions - 3)
        if args.t_train > _data.shape[-2]:
            args.t_train = _data.shape[-2]


        # collate_fn = partial(unet_collate, args=args, device=device)
        collate_fn = partial(unet_collate, args=args)
        # Build out dataloaders with pre processing 
        train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
                                  num_workers=0, shuffle=True, collate_fn=collate_fn)
        test_loader = DataLoader(test_dataset, batch_size=args.batch_size,
                                 num_workers=0, shuffle=False, collate_fn=collate_fn)
        
        return train_loader, test_loader
    
    elif args.model == "schnet" or args.model == "cgcnn":
        train_dataset = MD17SingleDataset("trajectory", args.molecule, "train", args.split, root="./data/MD17")
        test_dataset = MD17SingleDataset("trajectory", args.molecule, "test", args.split, root="./data/MD17")
        
        if args.dataset_length is not None:
            train_dataset = torch.utils.data.random_split(train_dataset, [args.dataset_length, len(train_dataset) - args.dataset_length])[0]
            test_dataset = torch.utils.data.random_split(test_dataset, [int(args.dataset_length/10), len(test_dataset) - int(args.dataset_length/10)])[0]
            print("truncated train_dataset", len(train_dataset))
        
        train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,
                                  collate_fn=qchem_collate, num_workers=0)

        test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False,
                                 collate_fn=qchem_collate, num_workers=0)
    
    elif args.model == "cnn" or args.model == "resnet":
        transform = transforms.Compose([
            transforms.Grayscale(num_output_channels=3),  # Convert to RGB format
            # transforms.Resize((224, 224)),  # Resize to match ViT's expected input size
            transforms.ToTensor(),
        ])
        train_dataset = datasets.MNIST(root = "./data/vision/",train = True,transform = ToTensor(),
            download = True)
        test_dataset = datasets.MNIST(root = "./data/vision/", train = False, transform = ToTensor(),
            download=True)
        
        if args.dataset_length is not None:
            train_dataset = torch.utils.data.random_split(train_dataset, [args.dataset_length, len(train_dataset) - args.dataset_length])[0]
            test_dataset = torch.utils.data.random_split(test_dataset, [int(args.dataset_length/10), len(test_dataset) - int(args.dataset_length/10)])[0]
            print("truncated train_dataset", len(train_dataset))
        
        train_loader= DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, 
                                 num_workers=0,collate_fn=vision_collate)
        test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0,
                                 collate_fn=vision_collate)
    
    elif args.model == "transformer" or args.model == "transformer2":
        transform = transforms.Compose([
            transforms.Grayscale(num_output_channels=3),  # Convert to RGB format
            # transforms.Resize((224, 224)),  # Resize to match ViT's expected input size
            transforms.ToTensor(),
        ])
        train_dataset = datasets.MNIST(root = "./data/vision/",train = True,transform = transform,
            download = False)
        test_dataset = datasets.MNIST(root = "./data/vision/", train = False, transform = transform,
            download=False)
        
        if args.dataset_length is not None:
            train_dataset = torch.utils.data.random_split(train_dataset, [args.dataset_length, len(train_dataset) - args.dataset_length])[0]
            # test_dataset = torch.utils.data.random_split(test_dataset, [int(args.dataset_length/10), len(test_dataset) - int(args.dataset_length/10)])[0]
            print("truncated train_dataset", len(train_dataset))

        train_loader= DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
        test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)
    
    return train_loader, test_loader


# =============================================================================
# Dataloader selection
# =============================================================================

def get_argparser(parser):
    args, _ = parser.parse_known_args()    
    if args.model == "unet":
        # [Unet params]
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=UNet1dWrap)
        parser.add_argument("-rb", "--reduced_batch", type=int, default=4)
        parser.add_argument("-rr", "--reduced_resolution", type=int, default=4)
        parser.add_argument("-rrt", "--reduced_resolution_t", type=int, default=5)
        parser.add_argument("-is", "--initial_step", type=int, default=10)
        parser.add_argument("-tt", "--t_train", type=str, default=200)
        parser.add_argument("-nc", "--num_channels", type=int, default=1)
        parser.add_argument("-wi", "--in_channels", type=int, default=1)
        parser.add_argument("-mo", "--out_channels", type=int, default=1)
        parser.add_argument("-bs", "--batch_size", type=int, default=50)
        parser.add_argument("-us", "--unroll_step", type=int, default=20)
        parser.add_argument("-cp", "--cloud_path", action='store_true')
    
    elif args.model == "schnet":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=SchNetWrap)
        parser.add_argument("-P", "--split", type=int, default=1000,
                            help="the name of dataset subset, aka the number of train samples")
        parser.add_argument("-m", "--molecule", default="a", type=str,
                            help="lowercase initial of the molecule in the dataset")
        parser.add_argument("-F", "--force", type=bool, default=True,
                            help="train with force")
        parser.add_argument("-b", "--batch_size", type=int, default=20,
                            help="batch size to train") 
    
    elif args.model == "cgcnn":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=CGCNN)
        parser.add_argument("-P", "--split", type=int, default=1000,
                            help="the name of dataset subset, aka the number of train samples")
        parser.add_argument("-m", "--molecule", default="a", type=str,
                            help="lowercase initial of the molecule in the dataset")
        parser.add_argument("-F", "--force", type=bool, default=True,
                            help="train with force")
        parser.add_argument("-b", "--batch_size", type=int, default=20,
                            help="batch size to train") 
    
    elif args.model == "resnet":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=ModifiedResNet18)
        parser.add_argument("-b", "--batch_size", type=int, default=128, help="batch size to train")
    
    elif args.model == "cnn":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=CNN)
        parser.add_argument("-b", "--batch_size", type=int, default=128, help="batch size to train")
    
    elif args.model == "transformer":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=Modifiedvit_b_16)
        parser.add_argument("-b", "--batch_size", type=int, default=128, help="batch size to train")
    
    elif args.model == "transformer2":
        parser.add_argument("-nm", "--nn_module", type=nn.Module, default=Modifiedvit_b_16_adjust)
        parser.add_argument("-b", "--batch_size", type=int, default=128, help="batch size to train")
        parser.add_argument("-num_heads", "--num_heads", type=int, default=8)
        parser.add_argument("-num_layers", "--num_layers", type=int, default=16)
        parser.add_argument("-mlp_dim", "--mlp_dim", type=int, default=1280)
        parser.add_argument("-hidden_dim", "--hidden_dim", type=int, default=320)
    
    if args.cache_size is None:
        args.cache_size = args.num_particles

    if args.view_size is None:
        args.view_size = args.cache_size

    return parser.parse_args()


def training_params(parser):
    args = get_argparser(parser)
    train_loader, test_loader = get_dataloaders(args)
    loss_fn = get_loss_fn(args)
    return args, train_loader, test_loader, loss_fn


# =============================================================================
# Wandb init
# =============================================================================

def wandb_init(args, dataloader):
    if args.wandb:
        timezone = pytz.timezone("America/Los_Angeles")
        start_dt = (datetime.now().astimezone(timezone)).strftime("%Y-%m-%d_%H-%M-%S")
        config = {
            "model": args.model,
            "train": args.train,
            "num_device": args.devices,
            "num_particles": args.num_particles,
            "cache_size": args.cache_size,
            "view_size": args.view_size,
            "learning_rate": args.learning_rate,
            "weight_decay": args.weight_decay,
            "batch_size": args.batch_size,
            "dataset size": len(dataloader.dataset),
        }
        if "num_params" in args:
            config["num_params"] = args.num_params
        print("Dataset size", len(dataloader.dataset))
        if args.train == args.train == "ensemble" or "ensemble_push":
            pass
        elif args.train == "mswag" or args.train == "mswag_push":
            config["pretrain_epochs"] = args.pretrain_epochs
            config["swag_epochs"] = args.swag_epochs
            config["num_models"] = args.num_models
            config["samples"] = args.samples
            config["scale"] = args.scale
            config["var_clamp"] = args.var_clamp
        elif args.train == "svgd" or args.train == "svgd_push":
            config["bandwidth"] = args.bandwidth

        wandb.init(
            project="scale_master", 
            entity="bogp", 
            group=args.group,
            name= f"{args.model}_{args.train}_{args.num_particles}_particles_{args.devices}_devices",
            id=start_dt, 
            config=config)
        

class MyTimer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, type, value, traceback):
        self.time = time.perf_counter() - self.start
        
experiments/train_util.py
lbai-lab/push
"""
particle.py
pfuture.py
push.py
"""
__version__ = "0.1.0"

push/__init__.py
lbai-lab/push

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from typing import *

from push.bayes.infer import Infer
from push.particle import Particle


# =============================================================================
# Helper
# =============================================================================

def mk_optim(params):
    """returns Adam optimizer"""
    # Limitiation must be global
    return torch.optim.Adam(params, lr=1e-5, weight_decay=1e-3)


# =============================================================================
# Deep Ensemble
# =============================================================================

def _deep_ensemble_main(particle: Particle, dataloader, loss_fn, epochs) -> None:
    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))
    # Training loop
    for e in tqdm(range(epochs)):
        losses = []
        for data, label in dataloader:
            loss = particle.step(loss_fn, data, label).wait()
            losses += [loss]
            for pid in other_particles:
                particle.send(pid, "ENSEMBLE_STEP", loss_fn, data, label)
        print(f"Average loss {particle.pid}", torch.mean(torch.tensor(losses)))


def _ensemble_step(particle: Particle, loss_fn, data, label, *args) -> None:
    particle.step(loss_fn, data, label, *args)


class Ensemble(Infer):
    """Ensemble class"""
    def __init__(self, mk_nn: Callable, *args: any, num_devices=1, cache_size=4, view_size=4) -> None:
        super(Ensemble, self).__init__(mk_nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size)
        
    def bayes_infer(self,
                    dataloader: DataLoader, epochs: int,
                    loss_fn=torch.nn.MSELoss(),
                    num_ensembles=2, mk_optim=mk_optim,
                    ensemble_entry=_deep_ensemble_main, ensemble_state={}, f_save=False):
        # 1. Create particles
        pids = [
            self.push_dist.p_create(mk_optim, device=(0 % self.num_devices), receive={
                "ENSEMBLE_MAIN": ensemble_entry
            }, state=ensemble_state)]
        for n in range(1, num_ensembles):
            pids += [self.push_dist.p_create(mk_optim, device=(n % self.num_devices), receive={
                "ENSEMBLE_STEP": _ensemble_step,
            }, state={})]

        # 2. Perform independent training
        self.push_dist.p_wait([self.push_dist.p_launch(0, "ENSEMBLE_MAIN", dataloader, loss_fn, epochs)])

        if f_save:
            self.push_dist.save()

# =============================================================================
# Deep Ensemble Training
# =============================================================================

def train_deep_ensemble(dataloader: Callable, loss_fn: Callable, epochs: int,
                        nn: Callable, *args, num_devices=1, cache_size=4, view_size=4,
                        num_ensembles=2, mk_optim=mk_optim,
                        ensemble_entry=_deep_ensemble_main, ensemble_state={}) -> None:
    ensemble = Ensemble(nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size)
    ensemble.bayes_infer(dataloader, epochs, loss_fn=loss_fn, num_ensembles=num_ensembles, mk_optim=mk_optim,
                         ensemble_entry=ensemble_entry, ensemble_state=ensemble_state)
    return ensemble.p_parameters()

push/bayes/ensemble.py
lbai-lab/push
import atexit
import torch
from torch.utils.data import DataLoader
from typing import *

import push.push as ppush


class Infer:
    """base Infer class
    """    
    def __init__(self, mk_nn: Callable, *args: any, num_devices=1, cache_size=4, view_size=4) -> None:
        self.mk_nn = mk_nn
        self.args = args
        self.num_devices = num_devices
        self.cache_size = cache_size
        self.view_size = view_size
        
        # Create a PusH Distribution
        self.push_dist = ppush.PusH(self.mk_nn, *self.args, cache_size=self.cache_size, view_size=self.view_size)
        atexit.register(self._cleanup)

    def bayes_infer(self, dataloader: DataLoader, epochs: int, **kwargs) -> None:
        """inference function"""
        raise NotImplementedError
    
    def p_parameters(self) -> List[List[torch.Tensor]]:
        """returns parameters"""
        return [self.push_dist.p_parameters(pid) for pid in self.push_dist.particle_ids()]

    def _cleanup(self):
        self.push_dist._cleanup()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.push_dist.__exit__(exc_type, exc_value, traceback)

push/bayes/infer.py
lbai-lab/push
import torch
from torch.distributions.normal import Normal
from torch.utils.data import DataLoader
from tqdm import tqdm
from typing import *


from push.bayes.infer import Infer
from push.particle import Particle
from push.bayes.utils import flatten, unflatten_like


# =============================================================================
# Helper
# =============================================================================

def mk_empty_optim(params):
    # Limitiation must be global
    """returns empty optimizer"""
    return None


# -----------------------------------------------------
# Prior
# -----------------------------------------------------
    
def normal_prior(params: Iterable[torch.Tensor]) -> list[torch.Tensor]:
    """returns grads"""
    normal = Normal(0.0, 1.0)
    grads = []
    for param in params:
        # TODO: change me for different NNs
        y = torch.sum(normal.log_prob(param.requires_grad_()))
        grads += [torch.autograd.grad(y, param, grad_outputs=torch.ones_like(y), create_graph=True,)[0]]
    return grads


# -----------------------------------------------------
# SVGD Kernels
# -----------------------------------------------------

def torch_squared_exp_kernel(x: torch.Tensor, y: torch.Tensor, length_scale: float) -> torch.Tensor:
    """torch  squared exp kernal"""
    diff = (x - y) / length_scale
    radius2 = torch.dot(diff, diff)
    return torch.exp(-0.5 * radius2)

def torch_squared_exp_kernel_grad(x: torch.Tensor, y: torch.Tensor, length_scale: float) -> torch.Tensor:
    """torch squared exp kernal grad"""
    prefactor = (x - y) / (length_scale ** 2)
    return -prefactor * torch_squared_exp_kernel(x, y, length_scale)


# =============================================================================
# SVGD
# =============================================================================

def _svgd_leader(particle: Particle, prior, loss_fn: Callable, lengthscale, lr, dataloader: DataLoader, epochs) -> None:
    # TODO comment explaining this function
    n = len(particle.particle_ids())
    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))

    for e in tqdm(range(epochs)):
        losses = []
        for data, label in dataloader:
            # 1. Step every particle
            fut = particle.step(loss_fn, data, label)
            futs = [particle.send(pid, "SVGD_STEP", loss_fn, data, label) for pid in other_particles]
            fut.wait(); [f.wait() for f in futs]

            # 2. Gather every other particles parameters
            particles = {pid: (particle.get(pid) if pid != particle.pid else 
                list(particle.module.parameters())) for pid in particle.particle_ids()}
            for pid in other_particles:
                particles[pid] = particles[pid].wait()

            # 3. Compute kernel and kernel gradients
            update = {}
            for pid1, params1 in particles.items():
                params1 = list(particles[pid1].view().parameters()) if pid1 != particle.pid else params1
                p_i = flatten(params1)
                update[pid1] = torch.zeros_like(p_i)
                for pid2, params2 in particles.items():
                    params2 = list(particles[pid2].view().parameters()) if pid2 != particle.pid else params2

                    # Compute kernel
                    p_j = flatten(params2)
                    p_j_grad = flatten([p.grad if p.grad is not None else torch.zeros_like(p) for p in params2])
                    update[pid1] += torch_squared_exp_kernel(p_j, p_i, lengthscale) * p_j_grad
                    update[pid1] += torch_squared_exp_kernel_grad(p_j, p_i, lengthscale)
                update[pid1] = update[pid1] / n

            # 4. Send kernel
            futs = [particle.send(pid, "SVGD_FOLLOW", lr, update[pid]) for pid in other_particles]
            [f.wait() for f in futs]
            _svgd_follow(particle, lr, update[particle.pid])

            loss = loss_fn(particle.forward(data).wait().to("cpu"), label)
            losses += [torch.mean(torch.tensor(loss))]

        print(f"Average loss {torch.mean(torch.tensor(losses))}")


def _svgd_leader_memeff(particle: Particle, prior, loss_fn: Callable, lengthscale, lr, dataloader: DataLoader, epochs) -> None:
    # TODO comment explaining this function
    n = len(particle.particle_ids())
    other_particles = list(filter(lambda x: x != particle.pid, particle.particle_ids()))

    for e in tqdm(range(epochs)):
        losses = []
        for data, label in dataloader:
            # 1. Step every particle
            fut = particle.step(loss_fn, data, label)
            futs = [particle.send(pid, "SVGD_STEP", loss_fn, data, label) for pid in other_particles]
            fut.wait(); [f.wait() for f in futs]

            # 2. Gather every other particles parameters
            particles = {pid: (particle.get(pid) if pid != particle.pid else 
                list(particle.module.parameters())) for pid in particle.particle_ids()}
            for pid in other_particles:
                particles[pid] = particles[pid].wait()

            def compute_update(pid1, params1):
                params1 = list(particles[pid1].view().parameters()) if pid1 != particle.pid else params1
                p_i = flatten(params1)
                update = torch.zeros_like(p_i)
                for pid2, params2 in particles.items():
                    params2 = list(particles[pid2].view().parameters()) if pid2 != particle.pid else params2

                    # Compute kernel
                    p_j = flatten(params2)
                    p_j_grad = flatten([p.grad if p.grad is not None else torch.zeros_like(p) for p in params2])
                    diff = (p_j - p_i) / lengthscale
                    radius2 = torch.dot(diff, diff)
                    k = torch.exp(-0.5 * radius2).item()
                    diff.mul_(-k / lengthscale)
                    update.add_(p_j_grad, alpha=k)
                    update.add_(diff)
                    # update += torch_squared_exp_kernel(p_j, p_i, lengthscale) * p_j_grad
                    # update += torch_squared_exp_kernel_grad(p_j, p_i, lengthscale)
                update = update / n
                return update

            # 3. Compute kernel and kernel gradients
            for pid1, params1 in particles.items():
                # 4. Send kernel
                if pid1 != particle.pid:
                    update = compute_update(pid1, params1)
                    particle.send(pid, "SVGD_FOLLOW", lr, update).wait()
            update = compute_update(particle.pid, particles[particle.pid])
            _svgd_follow(particle, lr, update)
            if particle.state["args"].model == "fno":
                (_data, grid) = data
                loss = loss_fn(particle.forward(_data, grid ,label).wait().to("cpu"), label)
            else:
                loss = loss_fn(particle.forward(data).wait().to("cpu"), label)

            losses += [torch.mean(torch.tensor(loss))]

        print(f"Average loss {torch.mean(torch.tensor(losses))}")


def _svgd_step(particle: Particle, loss_fn: Callable, data: torch.Tensor, label: torch.Tensor) -> None:
    # TODO comment explaining this function
    return particle.step(loss_fn, data, label)


def _svgd_follow(particle: Particle, lr: float, update: List[torch.Tensor]) -> None:
    # TODO comment explaining this function
    # 1. Unflatten
    params = list(particle.module.parameters())
    updates = unflatten_like(update.unsqueeze(0), params)
    
    # 2. Apply the update to the input particle
    with torch.no_grad():
        for p, up in zip(params, updates):
            p.add_(up.real, alpha=-lr)


class SteinVGD(Infer):
    """SteinVGD Class"""
    def __init__(self, mk_nn: Callable, *args: any, num_devices=1, cache_size=4, view_size=4) -> None:
        super(SteinVGD, self).__init__(mk_nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size)
        
    def bayes_infer(self,
                    dataloader: DataLoader, epochs: int,
                    prior=None, loss_fn=torch.nn.MSELoss(),
                    num_particles=1, lengthscale=1.0, lr=1e-3,
                    svgd_entry=_svgd_leader, svgd_state={}):
        pid_leader = self.push_dist.p_create(mk_empty_optim, device=0, receive={
            "SVGD_LEADER": svgd_entry
        }, state=svgd_state)
        pids = [pid_leader]
        for p in range(num_particles - 1):
            pid = self.push_dist.p_create(mk_empty_optim, device=((p + 1) % self.num_devices), receive={
                "SVGD_STEP": _svgd_step,
                "SVGD_FOLLOW": _svgd_follow
            }, state=svgd_state)
            pids += [pid]
        self.push_dist.p_wait([self.push_dist.p_launch(0, "SVGD_LEADER", prior, loss_fn, lengthscale, lr, dataloader, epochs)])


# =============================================================================
# SVGD Train
# =============================================================================

def train_svgd(dataloader: DataLoader, loss_fn: Callable, epochs: int, num_particles: int, nn: Callable, *args,
               lengthscale=1.0, lr=1e3, prior=None,
               num_devices=1, cache_size=4, view_size=4,
               svgd_entry=_svgd_leader, svgd_state={}) -> None:
    """SVGD training function"""
    with SteinVGD(nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size) as stein_vgd:
        stein_vgd.bayes_infer(dataloader, epochs,
                              prior=prior, loss_fn=loss_fn,
                              num_particles=num_particles, lengthscale=lengthscale, lr=lr,
                              svgd_entry=svgd_entry, svgd_state=svgd_state)
        return stein_vgd.p_parameters()

push/bayes/stein_vgd.py
lbai-lab/push
from torch.utils.data import DataLoader
from tqdm import tqdm
from typing import *
import torch

from push.bayes.infer import Infer
from push.particle import Particle
from push.bayes.utils import flatten, unflatten_like
from push.lib.utils import detach_to_cpu


# =============================================================================
# Helper
# =============================================================================

def mk_optim(params):
    """returns Adam optimizer"""
    return torch.optim.Adam(params, lr=1e-5, weight_decay=1e-3)


# =============================================================================
# Swag Training
# =============================================================================

def _swag_step(particle: Particle,
               loss_fn: Callable,
               data: torch.Tensor,
               label: torch.Tensor,
               *args: any) -> None:
    # calls one swag particle's step function
    particle.step(loss_fn, data, label,*args)


def update_theta(state, state_sq, param, param_sq, n):
    """update_theta"""
    for st, st_sq, p, p_sq in zip(state, state_sq, param, param_sq):
        st.data = (st.data * n + p.data)/(n+1)
        st_sq.data = (st_sq.data * n + p_sq.data)/(n+1)


def _swag_swag(particle: Particle, reset: bool) -> None:    
    # if reset, initializes mom1 and mom2, else updates mom1 and mom2
    state = particle.state
    if reset:
        state[particle.pid] = {
            "mom1": [param for param in particle.module.parameters()],
            "mom2": [param*param for param in particle.module.parameters()]
        }
    else:
        params = [param for param in particle.module.parameters()]
        params_sq = [param*param for param in particle.module.parameters()]
        update_theta(state[particle.pid]["mom1"], state[particle.pid]["mom2"], params, params_sq, state["n"])
        state["n"] += 1


def _mswag_particle(particle: Particle, dataloader, loss_fn: Callable,
                    pretrain_epochs: int, swag_epochs: int, swag_pids: list[int]) -> None:
    # training function for mswag particle
    other_pids = [pid for pid in swag_pids if pid != particle.pid]
    
    # Pre-training loop
    for e in tqdm(range(pretrain_epochs)):
        losses = []
        for data, label in dataloader:
            fut = particle.step(loss_fn, data, label)
            futs = [particle.send(pid, "SWAG_STEP", loss_fn, data, label) for pid in other_pids]
            losses += [fut.wait()]
        print("Average epoch loss", torch.mean(torch.tensor(losses)))
    
    # Initialize swag
    [particle.send(pid, "SWAG_SWAG", True) for pid in other_pids]
    _swag_swag(particle, True)
    
    # Swag epochs
    for e in tqdm(range(swag_epochs)):
        losses = []
        for data, label in dataloader:
            # Update
            futs = [particle.send(pid, "SWAG_STEP", loss_fn, data, label) for pid in other_pids]
            fut = particle.step(loss_fn, data, label)
            [f.wait() for f in futs]
            losses += [fut.wait()]
        futs = [particle.send(pid, "SWAG_SWAG", False) for pid in other_pids]
        _swag_swag(particle, False)
        [f.wait() for f in futs]
        print("Average epoch loss", torch.mean(torch.tensor(losses)))


# =============================================================================
# SWAG Inference
# =============================================================================

def _mswag_sample_entry(particle: Particle,
                        dataloader: DataLoader,
                        loss_fn: Callable,
                        scale: float,
                        var_clamp: float,
                        num_samples: int,
                        num_models) -> None:
    # TODO comment explaining this function
    # Unpack state
    state = particle.state
    pid = particle.pid
    other_pids = list(range(1, num_models))

    # Perform SWAG sampling for every particle
    my_ans = _mswag_sample(particle, dataloader, loss_fn, scale, var_clamp, num_samples)
    futs = [particle.send(pid, "SWAG_SAMPLE", dataloader, loss_fn, scale, var_clamp, num_samples) for pid in other_pids]
    ans = [f.wait() for f in futs]
    
    # Majority vote of majority prediction
    classes = {k: [0 for i in range(10)] for k in range(10)}
    max_preds = [my_ans['max_preds']] + [ans[i-1]['max_preds'] for i in range(1, num_models)]
    for idx, (data, label) in enumerate(dataloader):
        max_pred = torch.mode(torch.stack([max_preds[i][idx] for i in range(num_models)]), dim=0).values
        for x, y in zip(max_pred, label):
            classes[y.item()][x.item()] += 1
    
    # Majority vote across all SWAG samples
    all_preds = [my_ans['preds']] + [ans[i-1]['preds'] for i in range(1, num_models)]
    classes2 = {k: [0 for i in range(10)] for k in range(10)}
    for idx, (data, label) in enumerate(dataloader):
        preds = []
        for m in range(num_models):
            preds += [torch.stack([all_preds[m][i][idx] for i in range(num_samples)])]
        max_pred = torch.mode(torch.cat(preds, dim=0), dim=0).values
        for x, y in zip(max_pred, label):
            classes2[y.item()][x.item()] += 1

    return classes, classes2


def _mswag_sample(particle: Particle,
                  dataloader: DataLoader,
                  loss_fn: Callable,
                  scale: float,
                  var_clamp: float,
                  num_samples: int) -> None:
    """Inspired by https://github.com/wjmaddox/swa_gaussian/blob/master/swag/posteriors/swag.py
    """
    state = particle.state
    pid = particle.pid
    # Gather
    mean_list = [param for param in particle.state[pid]["mom1"]]
    sq_mean_list = [param for param in particle.state[pid]["mom2"]]

    scale_sqrt = scale ** 0.5
    mean = flatten(mean_list)
    sq_mean = flatten(sq_mean_list)

    # Compute original loss
    classes = {k: [0 for i in range(10)] for k in range(10)}
    losses = []
    for data, label in tqdm(dataloader):
        pred = detach_to_cpu(particle.forward(data).wait())
        loss = loss_fn(pred, label)
        cls = pred.softmax(dim=1).argmax(dim=1)
        for x, y in zip(cls, label):
            classes[y.item()][x.item()] += 1
        losses += [loss.detach().to("cpu")]

    # Compute and store prediction for each SWAG sample
    preds = {i: [] for i in range(num_samples)}
    swag_losses = {}
    for i in range(num_samples):
        # Draw diagonal variance sample
        var = torch.clamp(sq_mean - mean ** 2, var_clamp)
        var_sample = var.sqrt() * torch.randn_like(var, requires_grad=False)

        rand_sample = var_sample

        # Update sample with mean and scale
        sample = mean + scale_sqrt * rand_sample
        sample = sample.unsqueeze(0)

        # Update
        samples_list = unflatten_like(sample, mean_list)

        for param, sample in zip(particle.module.parameters(), samples_list):
            param.data = sample

        swag_losses = []
        for data, label in tqdm(dataloader):
            pred = detach_to_cpu(particle.forward(data).wait())
            cls = pred.softmax(dim=1).argmax(dim=1)
            preds[i] += [cls]
            loss = loss_fn(pred, label)
            swag_losses += [loss.detach().to("cpu")]    
    
    # Majority prediction for the current SWAG particle
    max_preds = []
    for n in range(len(preds[0])):
        max_preds += [torch.mode(torch.stack([preds[i][n] for i in range(num_samples)]), dim=0).values]

    return {
        "classes": classes, 
        "losses": torch.mean(torch.tensor(losses)),
        "preds": preds,
        "max_preds": max_preds,
    }


class MultiSWAG(Infer):
    """MultiSWAG class"""
    def __init__(self, mk_nn: Callable, *args: any, num_devices=1, cache_size=4, view_size=4) -> None:
        """init"""
        super(MultiSWAG, self).__init__(mk_nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size)
        self.swag_pids = []
        self.sample_pids = []
        
    def bayes_infer(self,
                    dataloader: DataLoader, 
                    loss_fn=torch.nn.MSELoss(),
                    num_models=1, lr=1e-3, pretrain_epochs=10, swag_epochs=5,
                    mswag_entry=_mswag_particle, mswag_state={}, f_save=False,
                    mswag_sample_entry=_mswag_sample_entry, mswag_sample=_mswag_sample):
        """bayes_infer"""
        if "n" in mswag_state:
            raise ValueError(f"Cannot run with state {mswag_state['n']}. Please rename.")
        mswag_state["n"] = 1

        def mk_swag(model_num):
            # Particle for parameter
            if model_num == 0:
                param_pid = self.push_dist.p_create(mk_optim, device=(model_num % self.num_devices), receive={
                    "SWAG_PARTICLE": mswag_entry,
                    "SWAG_SAMPLE_ENTRY": mswag_sample_entry,
                }, state=mswag_state)
            else:
                param_pid = self.push_dist.p_create(mk_optim, device=(model_num % self.num_devices), receive={
                    "SWAG_STEP": _swag_step,
                    "SWAG_SWAG": _swag_swag,
                    "SWAG_SAMPLE": mswag_sample
                }, state=mswag_state)
            return param_pid

        self.swag_pids = [mk_swag(n) for n in range(num_models)]
        for pid in self.swag_pids:
            if pid in mswag_state:
                raise ValueError(f"Cannot run with state {pid}. Please rename.")

        self.push_dist.p_wait([self.push_dist.p_launch(self.swag_pids[0], "SWAG_PARTICLE", dataloader, loss_fn, pretrain_epochs, swag_epochs, self.swag_pids)])

        if f_save:
            self.push_dist.save()

    def posterior_pred(self, dataloader: DataLoader, loss_fn=torch.nn.MSELoss(),
                       num_samples=20, scale=1.0, var_clamp=1e-30):
        """posterior_pred"""
        self.push_dist.p_wait([self.push_dist.p_launch(0, "SWAG_SAMPLE_ENTRY", dataloader, loss_fn, scale, var_clamp, num_samples, len(self.swag_pids))])


# =============================================================================
# Multi-Swag Training
# =============================================================================

def train_mswag(dataloader: DataLoader,
                loss_fn: Callable,
                pretrain_epochs: int,
                swag_epochs: int,
                num_models: int,
                cache_size: int,
                view_size:int,
                nn: Callable, *args,
                num_devices=1, lr=1e-3,
                mswag_entry=_mswag_particle, mswag_state={}, f_save=False,
                mswag_sample_entry=_mswag_sample_entry,
                mswag_sample=_mswag_sample):
    """train_mswag"""
    mswag = MultiSWAG(nn, *args, num_devices=num_devices, cache_size=cache_size, view_size=view_size)
    mswag.bayes_infer(dataloader, loss_fn, num_models, lr=lr, pretrain_epochs=pretrain_epochs,
                      swag_epochs=swag_epochs, mswag_entry=mswag_entry, mswag_state=mswag_state,
                      f_save=f_save, mswag_sample_entry=mswag_sample_entry, mswag_sample=mswag_sample)
    return mswag

push/bayes/swag.py
lbai-lab/push
import torch


def flatten(lst):
    """Inspired by: https://github.com/wjmaddox/swa_gaussian/blob/master/swag/utils.py"""
    tmp = [i.contiguous().view(-1, 1) for i in lst]
    return torch.cat(tmp).view(-1)


def unflatten_like(vector, likeTensorList):
    """ Inspired by: https://github.com/wjmaddox/swa_gaussian/blob/master/swag/utils.py"""
    # Takes a flat torch.tensor and unflattens it to a list of torch.tensors
    #    shaped like likeTensorList
    outList = []
    i = 0
    for tensor in likeTensorList:
        # n = module._parameters[name].numel()
        n = tensor.numel()
        outList.append(vector[:, i : i + n].view(tensor.shape))
        i += n
    return outList

push/bayes/utils.py
lbai-lab/push
from typing import *
import threading

import torch
import torch.nn as nn


class ParticleCache:
    """Loads particles on and off the accelerator.
    """    
    def __init__(self, mk_module: Callable, args: List[any], cache_size: int, device: int, lock: threading.Lock) -> None:
        # Module
        self.mk_module = mk_module
        self.args = args

        # Device
        self._device = device           # the device id
        
        # Cache
        self._cache_size = int(cache_size)   # maximum cache size
        self._cache2pid = {}            # cache position -> pid
        self._active2pid = {}           # cache position -> (pid, module)
        self._next_pos = 0

        # All particles
        self._module_disk = {}          # pid -> module (cpu)
        self._particle_disk = {}        # pid -> path (disk)
        self._optim_cache = {}          # pid -> Optimizer
        self._mk_optims = {}            # pid -> closure
        self._pid2cache = {}            # pid -> cache position
        
        self._pinned = set()            # pinned pid
        self._lock = lock               # lock for pinned pid

    def save_to_disk(self, pid, path="./particles"):
        module = self.mk_module(*self.args)
        if pid in self._module_disk:
            params, _ = self._module_disk[pid]
            for p, param in zip(module.parameters(), params):
                p.data = param 
            torch.save(module.state_dict(), self._particle_disk[pid])
        else:
            c_idx = self._pid2cache[pid]
            module = self._active2pid[c_idx][1]
            torch.save(module.state_dict(), self._particle_disk[pid])
        
    def _save_w_grads(self, pid: int, module: nn.Module) -> None:
        params = []
        params_grad = []
        for param in module.parameters():
            params_grad += [param.grad.detach().to("cpu") if param.grad is not None else None]
            params += [param.detach().to("cpu")]
        self._module_disk[pid] = (params, params_grad)

    def _load_w_grads(self, pid: int, module: nn.Module) -> None:
        params, params_grad = self._module_disk[pid]
        for p, param, param_grad in zip(module.parameters(), params, params_grad):
            p.data = param
            p.grad = param_grad
            
    def create(self, pid: int, mk_optim: Callable) -> nn.Module:
        # Create module
        module = self.mk_module(*self.args)
        module = module.to(self._device)
        # Save new module to disk
        self._module_disk[pid] = (module.parameters(), [param.grad for param in module.parameters()])

        c_idx = self._next_pos % self._cache_size
        if c_idx in self._active2pid:
            # There is an active particle sharing the place that we want to write our new particle into
            active_pid, active_module = self._active2pid.pop(c_idx)
            
            # Save active module to disk
            self._save_w_grads(active_pid, active_module)
            
            # Write new module into existing module
            active_module.load_state_dict(module.state_dict())
            
            # Update
            self._active2pid[c_idx] = (pid, active_module)
            module = active_module
        else:
            self._cache2pid[c_idx] = set()
            self._active2pid[c_idx] = (pid, module)
        # module should be the one in self._active2pid

        self._cache2pid[c_idx].add(pid)
        self._pid2cache[pid] = c_idx

        # Put in cache
        self._particle_disk[pid] = f"particles/device{self._device}_particle{pid}.pth"
        self._mk_optims[pid] = mk_optim
        self._optim_cache[pid] = mk_optim(module.parameters())

        # Increment
        self._next_pos += 1
        
        return module

    def try_pin(self, pid: int) -> bool:
        c_idx = self._pid2cache[pid]
        with self._lock:
            if self._active2pid[c_idx][0] == pid:
                self._pinned.add(pid)
                return True
            else:
                return False

    def release(self, pid, thread):
        with self._lock:
            thread.join()
            if pid in self._pinned:
                self._pinned.remove(pid)

    def unpin(self, pid: int) -> None:
        with self._lock:
            if pid in self._pinned:
                self._pinned.remove(pid)

    def try_read(self, pid: int, pin=False, msg=None) -> nn.Module:
        # if msg is not None:
        #     print(msg)
        c_idx = self._pid2cache[pid]
        with self._lock:
            if self._active2pid[c_idx][0] == pid:
                # Return if's active
                if pin:
                    self._pinned.add(pid)
                return self._active2pid[c_idx][1]
            elif self._active2pid[c_idx][0] in self._pinned:
                # Return None if its pinned
                return None
            
            # Pin
            if pin:
                self._pinned.add(pid)

            # Prepare to swap
            active_pid, active_module = self._active2pid.pop(c_idx)

            # Save with gradients
            self._save_w_grads(active_pid, active_module)

            # Remove old optimizer
            old_optim = self._optim_cache.pop(active_pid)
            del old_optim 

            # Load particle into active module
            self._load_w_grads(pid, active_module)
            self._active2pid[c_idx] = (pid, active_module.to(self._device))            
            
            # Return new module
            new_module = self._active2pid[c_idx][1]
            
            # Restore optim
            self._optim_cache[pid] = self._mk_optims[pid](new_module.parameters())
            
            # Result
            return new_module

    def contains(self, pid):
        return pid in self._pid2cache

    def particles(self) -> List[int]:
        return self._pid2cache.keys()
    
    def __str__(self) -> str:
        return f"active2pid: {str({k: v[0] for k, v in self._active2pid.items()})}\n cache2pid: {str(self._cache2pid)}"
    

class ParticleCacheLRU:
    """Loads particles on and off the accelerator.
    """    
    def __init__(self, mk_module: Callable, args: List[any], cache_size: int, device: int) -> None:
        # Module
        self.mk_module = mk_module
        self.args = args

        # Device
        self._device = device           # the device id
        
        # Cache
        self._cache_size = cache_size   # maximum cache size
        self._module_cache = {}         # pid -> module (device)
        self._module_disk = {}          # pid -> module (cpu)
        self._particle_disk = {}        # pid -> path (disk)
        self._optim_cache = {}          # pid -> Optimizer
        self._mk_optims = {}            # pid -> closure
        self._lru = []

    def _save(self, pid: int, module: nn.Module, disk=False) -> None:
        if disk:
            torch.save(module.state_dict(), self._particle_disk[pid])
        else:
            tmp = self.mk_module(*self.args)
            tmp.load_state_dict(module.state_dict())
            self._module_disk[pid] = tmp

    def _load(self, pid: int, module: nn.Module, disk=False) -> None:
        if disk:
            checkpoint = torch.load(self._particle_disk[pid])
            module.load_state_dict(checkpoint)
        else:
            module.load_state_dict(self._module_disk[pid].state_dict())
            
    def read(self, pid: int) -> nn.Module:
        if pid in self._module_cache:
            return self._module_cache[pid]
        else:
            if len(self._module_cache) >= self._cache_size:
                # Remove particle
                lru_idx = self._lru.pop(0)
                module = self._module_cache.pop(lru_idx)
                self._save(lru_idx, module)
        
                # Remove old optimizer
                old_optim = self._optim_cache.pop(pid)
                del old_optim 

                # Load particle
                self._load(pid, module)
                self._module_cache[pid] = module.to(self._device)
                self._lru += [pid]
                new_module = self._module_cache[pid]
                
                # Restore optim
                self._optim_cache[pid] = self._mk_optims[pid](new_module.parameters())
                
                # Result
                return new_module
            else:
                raise ValueError("Shouldn't happen ...")

    def write(self, pid: int, module: nn.Module) -> None:
        if pid in self._module_cache:
            self._module_cache[pid] = module
        else:
            if len(self._module_cache) >= self._cache_size:
                lru_idx = self._lru.pop(0)
                module_p = self._module_cache.pop(lru_idx)
                self._save(lru_idx, module_p)

            self._module_cache[pid] = module.to(self._device)
            self._lru += [pid]

    def create(self, pid: int, mk_optim: Callable) -> nn.Module:
        self._particle_disk[pid] = f"particles/device{self._device}_particle{pid}.pth"
        module = self.mk_module(*self.args)
        module = module.to(self._device)
        self.write(pid, module)
        
        self._mk_optims[pid] = mk_optim
        self._optim_cache[pid] = mk_optim(module.parameters())

        return module

    def contains(self, pid):
        return pid in self._module_cache

    def particles(self) -> List[int]:
        return self._module_cache.keys()

push/lib/context_switch.py
lbai-lab/push
import torch
from typing import *


class MSG:
    pass


# =============================================================================
# Device Event Loop Messages
# =============================================================================

class NodeEvtLoopInitMSG(MSG):
    def __init__(self):
        pass


class NodeEvtLoopCleanupMSG(MSG):
    def __init__(self):
        pass


class NELBroadcastParticlesMSG(MSG):
    def __init__(self, in_queues, out_queues, particle_to_device):
        self.in_queues = in_queues
        self.out_queues = out_queues
        self.particle_to_device = particle_to_device


class NELBroadcastParticlesAckMSG(MSG):
    def __init__(self):
        pass


class NELSaveModel(MSG):
    def __init__(self, pid_fid):
        self.pid_fid = pid_fid


class NELSaveModelAckPDMSG(MSG):
    def __init__(self, pid_fid):
        self.pid_fid = pid_fid


# =============================================================================
# Push Distribution Messages
# =============================================================================

# -----------------------------------------------------
# One-Time
# -----------------------------------------------------

class ReceiveParticleInitPDMSG(MSG):
    def __init__(self, device: int, pid: int, mk_optim: Callable, receive, state):
        self.device = device
        self.pid = pid
        self.mk_optim = mk_optim
        self.receive = receive
        self.state = state


class ReceiveParticleInitAckPDMSG(MSG):
    def __init__(self):
        pass


class ReceiveRegisterPDMSG(MSG):
    def __init__(self, pid: int, msg: str, fn: Callable, state: Dict[str, any]):
        self.pid = pid
        self.msg = msg
        self.fn = fn
        self.state = state


class ReceiveRegisterAckPDMSG(MSG):
    def __init__(self):
        pass


# -----------------------------------------------------
# Multi-Time
# -----------------------------------------------------

class ReceiveFuncPDMSG(MSG):
    def __init__(self, pid_fid, pid_to: int, msg: str, args: List[any]):
        self.pid_fid = pid_fid
        self.pid_to = pid_to
        self.msg = msg
        self.args = args


class ReceiveFuncAckPDMSG(MSG):
    def __init__(self, pid_fid):
        self.pid_fid = pid_fid


class ReceiveParametersPDMSG(MSG):
    def __init__(self, pid_fid, pid: int):
        self.pid_fid = pid_fid
        self.pid = pid


class ReceiveParametersAckPDMSG(MSG):
    def __init__(self, pid_fid, params: List[torch.Tensor]):
        self.pid_fid = pid_fid
        self.params = params


# =============================================================================
# Particle Messages
# =============================================================================

class ReceiveFuncMSG(MSG):
    def __init__(self, pid_fid, pid: int, msg_name: str, args: List[any]):
        self.pid_fid = pid_fid
        self.pid = pid
        self.msg_name = msg_name
        self.args = args


class ReceiveFuncAckMSG(MSG):
    def __init__(self):
        pass


class ReceiveGetMSG(MSG):
    def __init__(self, pid_fid, pid_caller: int, pid: int):
        self.pid_fid = pid_fid
        self.pid_caller = pid_caller
        self.pid = pid


class ReceiveGetAckMSG(MSG):
    def __init__(self, fid: int, pid: int, params: List[torch.Tensor], params_grad: List[torch.Tensor]):
        self.fid = fid
        self.pid = pid
        self.params = params
        self.params_grad = params_grad

push/lib/messages.py
lbai-lab/push
from collections import OrderedDict
import threading
import torch
import torch.nn as nn
import torch.multiprocessing as mp
from typing import *

from push.lib.context_switch import ParticleCacheLRU, ParticleCache
from push.lib.messages import *
from push.lib.utils import detach_to_device, detach_to_cpu
from push.lib.waitable import Waitable
from push.particle import Particle, ParticleView
from push.pfuture import PFuture


class NodeEventLoop(Waitable):
    """
    1. Each node event loop is mapped to all devices on that node.
    2. The node event loop is responsible for executing all operations on particles.
    """    
    def __init__(self,
                 mk_module: Callable,
                 args: List[any],
                 in_queue: mp.Queue,
                 out_queue: mp.Queue,
                 rank: int,
                 devices: int,
                 cache_size: int,
                 view_size: int) -> None:
        # Node information
        self.rank = rank                         # Rank of NEL
        self.devices = devices                   # Devices on this NEL
        
        # Messaging back out to PusH
        self.in_queue = in_queue                 # receiving message queue
        self.out_queue = out_queue               # direct queue to PusH

        # Point-to-point messaging between particles
        self._in_queues = {}                     # Type: pid -> mp.Queue
        self._out_queues = {}                    # Type: pid -> mp.Queue

        # Particle information
        self._particle_to_rank = {}              # Type: pid -> rank
        self._particle_to_device = {}            # Type: pid -> device
        self._particle_to_state = {}             # Type: pid -> state
        self._hooks = {}                         # Type: pid -> msg -> closure, for triggering events

        # Device management
        self._device_to_pthread = OrderedDict()  # Type: device -> (pid, thread)
        self.particle_caches = {}                # Type: device -> cache, particle caches on this NEL
        self.view_caches = {}                    # Type: device -> cache, cross-device view caches on this NEL
        for device in devices:
            self.particle_caches[device] = ParticleCache(mk_module, args, cache_size, device, threading.Lock())
            self.view_caches[device] = ParticleCache(mk_module, args, view_size, device, threading.Lock())
        
        # Futures state for particles on this event loop
        self._future_id = 0                      # Type: fid, future id
        self._particle_to_futures = {}           # Type: pid -> set[fid]
        self._future_to_particle = {}            # Type: fid -> pid
        
        self._results = {}                       # Type: fid -> result
        self._results_lock = threading.Lock()

    # -----------------------------------------------------
    # Context switching
    # -----------------------------------------------------

    def _wait_particle_thread(self, pid: int) -> None:
        pid_device = self._particle_to_device[pid]
        if pid_device in self._device_to_pthread:
            # If device is in use, finish execution of active particle that is using that device
            active_pid, active_thread = self._device_to_pthread.pop(pid_device)
            active_pid_device = self._particle_to_device[active_pid]
            self.particle_caches[active_pid_device].release(active_pid, active_thread)

    def _context_switch_module(self, pid: int, pin=False, msg=None) -> nn.Module:
        pid_device = self._particle_to_device[pid]
        
        # Try read
        self._wait_particle_thread(pid)
        module = self.particle_caches[pid_device].try_read(pid, pin=pin, msg=msg)
        
        # Try reload module
        while module is None:
            # If module read fails, finish execution of active particle that is using that device
            active_pid, active_thread = self._device_to_pthread.popitem(last=False)
            active_pid_device = self._particle_to_device[active_pid]
            self.particle_caches[active_pid_device].release(active_pid, active_thread)
            
            # Try reload module
            module = self.particle_caches[pid_device].try_read(pid, pin=pin)
        return module

    # -----------------------------------------------------
    # Node event loop entry
    # -----------------------------------------------------

    def _start_event_loop(self) -> None:
        go = True
        while go:
            # Obtain message and dispatch
            msg = self.in_queue.get()
            go = self._dispatch(msg)

    def _dispatch(self, msg: MSG) -> bool:

        # -----------------------------------------------------
        # Particle functionality
        # -----------------------------------------------------
        
        if isinstance(msg, ReceiveFuncMSG):
            # Handle message
            sender, pid, msg_name = msg.pid_fid, msg.pid, msg.msg_name
            if msg_name in self._hooks[pid]:
                fn = self._hooks[pid][msg_name]
                state = self._particle_to_state[pid]
                pid_device = self._particle_to_device[pid]
                self._context_switch_module(pid, msg='ReceiveFuncMsg')
                particle = Particle(self, pid_device, pid, module, state)
                args = [detach_to_device(pid_device, arg) for arg in msg.args]
                fn(particle, *args)

        elif isinstance(msg, ReceiveGetMSG):
            # Handle message
            pid_device = self._particle_to_device[msg.pid]
            module = self._context_switch_module(pid)
            params = []
            params_grad = []
            for p in module.parameters():
                params += [p.detach().clone().cpu()]
                params_grad += [p.grad.detach().clone().cpu() if p.grad is not None else None]
            
            # Acknowledge
            caller_pid_device = self._particle_to_device[msg.pid_caller]
            self._in_queues[caller_pid_device].put(ReceiveGetAckMSG(msg.pid_fid, msg.pid, params, params_grad))

        # -----------------------------------------------------
        # PNN functionality
        # -----------------------------------------------------

        elif isinstance(msg, ReceiveParametersPDMSG):
            # Handle message
            pid_device = self._particle_to_device[msg.pid]
            module = self._context_switch_module(msg.pid)
            params = [x.detach().to("cpu") for x in module.parameters()]

            # Acknowledge
            self.out_queue.put(ReceiveParametersAckPDMSG(msg.pid_fid, params))

        elif isinstance(msg, ReceiveRegisterPDMSG):
            # Handle message
            self._hooks[msg.pid][msg.msg] = msg.fn

            # Acknowledge
            self.out_queue.put(ReceiveRegisterAckPDMSG())

        elif isinstance(msg, ReceiveParticleInitPDMSG):
            # Handle message
            self._particle_to_rank[msg.pid] = self.rank       # set rank
            self._particle_to_device[msg.pid] = msg.device    # set device
            module = self.particle_caches[self._particle_to_device[msg.pid]].create(msg.pid, msg.mk_optim)
            self._particle_to_state[msg.pid] = msg.state
            
            # Register receives
            self._hooks[msg.pid] = {}
            for recv_msg, fn in msg.receive.items():
                self._hooks[msg.pid][recv_msg] = fn
            self._particle_to_futures[msg.pid] = set()

            # Acknowledge
            self.out_queue.put(ReceiveParticleInitAckPDMSG())

        elif isinstance(msg, NELSaveModel):
            # Handle message
            for pid, pid_device in self._particle_to_device.items():
                self.particle_caches[pid_device].save_to_disk(pid)

            # Acknowledge
            self.out_queue.put(NELSaveModelAckPDMSG(msg.pid_fid))

        # -----------------------------------------------------
        # Node Event Loop Functionality
        # -----------------------------------------------------

        elif isinstance(msg, NELBroadcastParticlesMSG):
            # Handle message
            self._in_queues = msg.in_queues
            self._out_queues = msg.out_queues
            self._particle_to_device = msg.particle_to_device

            # Acknowledge
            self.out_queue.put(NELBroadcastParticlesAckMSG())

        elif isinstance(msg, NodeEvtLoopCleanupMSG):
            return False

        # -----------------------------------------------------
        # PNN - particle functionality
        # -----------------------------------------------------
        
        elif isinstance(msg, ReceiveFuncPDMSG):
            # Handle message
            sender, pid, msg_name = msg.pid_fid, msg.pid_to, msg.msg
            if msg_name in self._hooks[pid]:
                fn = self._hooks[pid][msg_name]
                state = self._particle_to_state[pid]
                pid_device = self._particle_to_device[pid]
                module = self._context_switch_module(pid)
                particle = Particle(self, pid_device, pid, module, state)
                try:
                    y = fn(particle, *msg.args)
                except Exception as e:
                    self.out_queue.put(e)
                    self._cleanup()
                    raise e
                
                # Acknowledge
                self.out_queue.put(ReceiveFuncAckPDMSG(sender))

        return True        

    def _dispatch_receive_get_ack(self, msg: ReceiveGetAckMSG) -> None:
        # Create particle on this device event loop if it doesn't exist
        pid_device = self._particle_to_device[msg.pid]
        if not self.view_caches[pid_device].contains(msg.pid):
            module = self.view_caches[pid_device].create(msg.pid, lambda x: None)
        else:
            module = self.view_caches[pid_device].read(msg.pid)

        # Copy parameters over
        for param, p, p_grad in zip(module.parameters(), msg.params, msg.params_grad):
            with torch.no_grad():
                param.copy_(p)
            if p_grad is not None:
                param.grad = p_grad.to(pid_device)

        # Return a lazy view
        self._results[msg.fid] = ParticleView(self.view_caches[self._particle_to_device[msg.pid]], msg.pid)

    def _wait(self, fid: int) -> any:
        def check_results(fid: int):
            # Check to see if we already have the result
            if fid in self._results:
                result = self._results.pop(fid)
                pid = self._future_to_particle.pop(fid)
                self._particle_to_futures[pid].remove(fid)
                return True, result
            else:
                return False, None

        result = check_results(fid)
        msgs = []
        while not result[0]:
            msg = self.in_queue.get()
            if isinstance(msg, ReceiveGetAckMSG):
                self._dispatch_receive_get_ack(msg)
                result = check_results(fid)
            elif isinstance(msg, ReceiveFuncMSG) or isinstance(msg, ReceiveFuncPDMSG):
                # Note: Handle this particles receive
                #       Additional parallelism here if we are not on the same PID
                msgs += [msg]
            else:
                # Dispatch PNN messages
                go = self._dispatch(msg)
        
        # Dispatch buffered messages
        for msg in msgs:
            go = self._dispatch(msg)
        
        return result[1]

    def _cleanup(self) -> None:
        self._dispatch(NodeEvtLoopCleanupMSG())

    # -----------------------------------------------------
    # Future functionality
    # -----------------------------------------------------

    def _create_future_id(self) -> int:
        fid = self._future_id
        self._future_id += 1
        return fid

    def _register_future(self, pid: int, fid: int) -> None:
        self._particle_to_futures[pid].add(fid)
        self._future_to_particle[fid] = pid

    # -----------------------------------------------------
    # Particle functionality
    # -----------------------------------------------------

    def particles(self) -> List[int]:
        return list(self._particle_to_device.keys())

    def register_receive(self, pid: int, msg:str, fn: Callable, state: dict) -> None:
        self._hooks[pid][msg] = (fn, state)

    def send(self, send_particle: Particle, pid_curr: int, pid: int, msg_name: str, *args: any) -> PFuture:
        # Create future
        fid = self._create_future_id()
        self._register_future(pid_curr, fid)

        # Check communication
        rank_id_curr = self._particle_to_rank[pid_curr]
        rank_id = self._particle_to_rank[pid]
        if rank_id == rank_id_curr: # We are on the same rank
            # NOTE: INVARIANT: the current function always has it's module in scope
            module = self._context_switch_module(pid, msg=f'send {msg_name} from {pid_curr} to {pid}')
            
            # NOTE: Execute function on main node event loop.
            #       Compute heavy items are run on separate threads.
            fn = self._hooks[pid][msg_name]
            state = self._particle_to_state[pid]
            pid_device = self._particle_to_device[pid]
            particle = Particle(self, pid_device, pid, module, state)
            args2 = [detach_to_device(pid_device, arg) for arg in args]
            try:
                y = fn(particle, *args2)
            except Exception as e:
                self.out_queue.put(e)
                self._cleanup()
                raise e
            self._results[fid] = y

            # NOTE: INVARIANT: the current function always has it's module in scope
            send_particle.module = self._context_switch_module(pid_curr, msg='send switch back')

            # Return future
            return PFuture(self, pid_curr, pid, fid)
        else:
            # Need to communicate
            args2 = [detach_to_cpu(arg) for arg in args]
            self._in_queues[rank_id].put(ReceiveFuncMSG((pid_curr, fid), pid, msg_name, args2))
            
            # Return future
            return PFuture(self, pid_curr, pid, fid)

    def get(self, pid_curr: int, pid: int) -> PFuture:
        # Create future
        fid = self._create_future_id()
        self._register_future(pid_curr, fid)

        # Check communication
        rank_id_curr = self._particle_to_rank[pid_curr]
        rank_id = self._particle_to_rank[pid]        
        if rank_id == rank_id_curr: # We are on the same rank
            # NOTE: INVARIANT: the current function always has it's module in scope
            module = self._context_switch_module(pid, msg='get')

            # Create space for pid on pid_curr's device
            device_curr = self._particle_to_device[pid_curr]
            if not self.view_caches[device_curr].contains(pid):
                module_on_curr = self.view_caches[device_curr].create(pid, lambda x: None)
            else:
                module_on_curr = self.view_caches[device_curr].try_read(pid)

            if True:
                # Copy parameters over
                for param, p in zip(module_on_curr.parameters(), module.parameters()):
                    with torch.no_grad():
                        param.copy_(p).to(device_curr)
                    if p.grad is not None:
                        param.grad = p.grad.to(device_curr)

                self._results[fid] = ParticleView(self.view_caches[device_curr], pid)
                return PFuture(self, pid_curr, pid, fid)
            else:
                def worker(device_curr, module_on_curr, module):
                    for param, p in zip(module_on_curr.parameters(), module.parameters()):
                        with torch.no_grad():
                            param.copy_(p).to(device_curr)
                        if p.grad is not None:
                            param.grad = p.grad.to(device_curr)

                        self._results[fid] = ParticleView(self.view_caches[device_curr], pid)

                # Launch thread
                t = threading.Thread(target=worker, args=(device_curr, module_on_curr, module,))
                self._device_to_pthread[device_curr] = (pid_curr, t)
                t.start()
                # Return future
                return PFuture(self, pid_curr, pid, fid, t=t)
        else:
            # Need to communicate
            self._in_queues[rank_id].put(ReceiveGetMSG(fid, pid_curr, pid))

            # Return future
            return PFuture(self, pid_curr, pid, fid)

    def wait(self, pfutures: List[PFuture]) -> List[any]:
        acc = []
        for pfuture in pfutures:
            acc += [pfuture.wait()]
        return acc

    # -----------------------------------------------------
    # NN Functionality
    # -----------------------------------------------------

    def zero_grad(self, pid: int) -> PFuture:
        # Register future
        fid = self._create_future_id()
        self._register_future(pid, fid)

        def worker(module):
            module.zero_grad()
            self._results[fid] = None

        # Context switch
        module = self._context_switch_module(pid, msg='zero_grad')

        # Launch thread
        t = threading.Thread(target=worker, args=(module,))
        pid_device = self._particle_to_device[pid]
        self._device_to_pthread[pid_device] = (pid, t)
        t.start()

        # Return future
        return PFuture(self, pid, pid, fid, t=t)

    def forward(self, pid: int, x: torch.Tensor, *args: any) -> PFuture:
        # Register future
        fid = self._create_future_id()
        self._register_future(pid, fid)

        # Functionality for forward
        def worker(module):
            try:
                y = module.forward(x, *args)
                self._results[fid] = y
            except Exception as e:
                self.out_queue.put(e)
                self._cleanup()
                raise e

        # Context switch
        module = self._context_switch_module(pid, msg='forward')

        # Launch thread
        t = threading.Thread(target=worker, args=(module,))
        pid_device = self._particle_to_device[pid]
        self._device_to_pthread[pid_device] = (pid, t)
        t.start()
        
        # Return future
        return PFuture(self, pid, pid, fid, t=t)

    def step(self, pid: int, loss_fn: Callable, data: torch.Tensor, label: torch.Tensor, *args: any) -> PFuture:
        # Register future
        fid = self._create_future_id()
        self._register_future(pid, fid)
        
        # Functionality for step
        def worker(module):
            try:
                # Backwards pass
                module.zero_grad()
                y = module.forward(data, *args)
                loss = loss_fn(y, label)
                loss.backward()

                # Optionally step
                pid_device = self._particle_to_device[pid]
                optim = self.particle_caches[pid_device]._optim_cache[pid]
                if optim:
                    optim.step()

                # NOTE: Assumes that dictionary stores are atomic
                self._results[fid] = loss
            except Exception as e:
                self.out_queue.put(e)
                self._cleanup()
                raise e

        # Context switch
        module = self._context_switch_module(pid, pin=True, msg=f'step {pid}')
        
        # Launch thread
        t = threading.Thread(target=worker, args=(module,))
        pid_device = self._particle_to_device[pid]
        self._device_to_pthread[pid_device] = (pid, t)
        t.start()

        # Return future
        return PFuture(self, pid, pid, fid, t=t)
    
push/lib/node_event_loop.py
lbai-lab/push
import torch
from typing import *


def to_device(device: int, val: Union[Dict, List, Tuple, torch.Tensor]) -> Union[Dict, List, torch.Tensor]:
    if isinstance(val, dict):
        return {k: to_device(device, v) for k, v in val.items()}
    elif isinstance(val, list):
        return [to_device(device, v) for v in val]
    elif isinstance(val, tuple):
        return tuple(to_device(device, v) for v in val)
    elif isinstance(val, torch.Tensor):
        return val.to(device)
    else:
        return val


def detach_to_device(device: int, val: Union[Dict, List, Tuple, torch.Tensor], requires_grad=False) -> Union[Dict, List, torch.Tensor]:
    if isinstance(val, dict):
        return {k: detach_to_device(device, v, requires_grad=requires_grad) for k, v in val.items()}
    elif isinstance(val, list):
        return [detach_to_device(device, v, requires_grad=requires_grad) for v in val]
    elif isinstance(val, tuple):
        return tuple(detach_to_device(device, v, requires_grad=requires_grad) for v in val)
    elif isinstance(val, torch.Tensor):
        return val.detach().requires_grad_().to(device) if requires_grad and isinstance(val, torch.FloatTensor) else val.detach().to(device)
    else:
        return val


def detach_to_cpu(val: Union[Dict, List, torch.Tensor]) -> Union[Dict, List, torch.Tensor]:
    if isinstance(val, dict):
        return {k: detach_to_cpu(v) for k, v in val.items()}
    elif isinstance(val, list):
        return [detach_to_cpu(v) for v in val]
    elif isinstance(val, torch.Tensor):
        return val.detach().to("cpu")
    else:
        return val

push/lib/utils.py
lbai-lab/push
from typing import *


class Waitable:
    def __init__(self):
        pass
    
    def _wait(self, fid: int) -> any:
        raise NotImplementedError
    
push/lib/waitable.py
lbai-lab/push
import torch
import torch.nn as nn
from typing import *

from push.lib.utils import detach_to_device
from push.lib.waitable import Waitable
from push.pfuture import PFuture


class Particle(Waitable):
    """User-facing particle interface.

    Implements concurrent particles.
    """
    def __init__(self, node_event_loop, device: int, pid: int, module: nn.Module, state: Dict[str, any]):
        assert isinstance(device, int)
        assert isinstance(pid, int)
        assert isinstance(module, nn.Module)

        self.device = device                     # Device the particle is on
        self.module = module                     # The module corresponding to the particle
        self.pid = pid                           # The particles identifier
        self.state = state                       # Particle's state

        # NOTE: WARNING: DO NOT ACCESS!!!
        self._node_event_loop = node_event_loop  # The device event loop that the particle is part of

    # -----------------------------------------------------
    # Particle functionality
    # -----------------------------------------------------

    def particle_ids(self) -> List[int]:
        """Returns all particles.

        Returns:
            List[int]: List of all particle identifiers visible to current particle.
        """
        return self._node_event_loop.particles()

    def register_receive(self, msg: str, fn: Callable, state: dict[str, any]) -> None:
        """Register receive functionality for current particle.

        Args:
            msg (str): Message for current particle to respond to.
            fn (Callable): Function to execute on `msg`.
            state (dict[str, any]): User state.

        Returns:
            None.
        """        
        return self._node_event_loop.register_receive(self.pid, msg, fn, state)

    def send(self, pid: int, msg: str, *args: any) -> PFuture:   
        """Send a `msg` to `pid` with arguments `*args` from current particle.

        Args:
            pid (int): Particle identifier of receiver.
            msg (str): Message.
            *args (any): Any additional arguments to give to receiver.

        Returns:
            PFuture: Handle to future result.
        """        
        return self._node_event_loop.send(self, self.pid, pid, msg, *args)

    def get(self, pid: int) -> PFuture:
        """Obtains particle `pid`'s parameters.

        Args:
            pid (int): Particle identifier to obtain parameters of.

        Returns:
            PFuture: Handle to future result.
        """        
        return self._node_event_loop.get(self.pid, pid)

    def wait(self, pfutures: List[PFuture]) -> List[any]:
        """Pause execution of current particle until all `pfutures` have been resolved.

        Args:
            pfutures (List[PFuture]): Future values to wait on.

        Returns:
            List[any]: List of resolved future values.
        """        
        return self._node_event_loop.wait(pfutures)

    # -----------------------------------------------------
    # NN functionality
    # -----------------------------------------------------

    def zero_grad(self) -> PFuture:
        """Zeros gradients on current particle.

        Returns:
            PFuture: Wait ensures that the gradients have been zeroed.
        """        
        self._node_event_loop.zero_grad(self.pid)

    def forward(self, x: torch.Tensor, *args: any) -> PFuture:
        """Performs a forward pass.

        Args:
            x (torch.Tensor): Input to the particle.

        Returns:
            PFuture: Future that eventually contains the value of the forward pass.
        """        
        x = detach_to_device(self.device, x, requires_grad=True)
        args = [detach_to_device(self.device, arg, requires_grad=True) for arg in args]
        # x = to_device(self.device, x, requires_grad=True)
        # args = [to_device(self.device, arg, requires_grad=True) for arg in args]
        return self._node_event_loop.forward(self.pid, x, *args)
        
    def step(self, loss_fn: Callable, data: torch.Tensor, label: torch.Tensor, *args) -> PFuture:
        """Performs a forward and backward pass using the registered optimizer.

        Args:
            loss_fn (Callable): Loss function to take a step with respect to.
            data (torch.Tensor): Data.
            label (torch.Tensor): label.

        Returns:
            PFuture: Future that eventually contains the loss of the step.
        """        
        data = detach_to_device(self.device, data)
        label = detach_to_device(self.device, label)
        args = [detach_to_device(self.device, arg) for arg in args]
        return self._node_event_loop.step(self.pid, loss_fn, data, label, *args)


class ParticleView:
    """User-facing particle view interface.

    Enables you to view another particle's parameters.
    """
    def __init__(self, view_cache, pid: int):
        self._view_cache = view_cache
        self.pid = pid

    def view(self) -> nn.Module:
        """Call to access it's latest state.

        Returns:
            nn.Module: The other particle.
        """        
        return self._view_cache.try_read(self.pid)

push/particle.py
lbai-lab/push
import threading
from typing import *

from push.lib.waitable import Waitable


class PFuture:
    """Particle Future. Eventually contains a result.
    """    
    def __init__(self, waitable: Waitable, send_pid: int, recv_pid: int, fid: int, t: Union[None, threading.Thread]=None):
        self._waitable = waitable   # waitable control flow
        self._send_pid = send_pid   # pid of particle that sends message
        self._recv_pid = recv_pid   # pid of particle that receives message and contains result
        self._fid = fid             # unique future identifier
        self._t_or_ts = t           # potential thread to block on

    def wait(self) -> any:
        """Wait until we get a result.

        Returns:
            any: The value that we waited to get.
        """        
        if self._t_or_ts is None:
            return self._waitable._wait(self._fid)
        elif isinstance(self._t_or_ts, threading.Thread):
            self._waitable._wait_particle_thread(self._recv_pid)
            return self._waitable._wait(self._fid)
        else:
            raise ValueError("Shouldn't happen ...")

    def __eq__(self, other) -> bool:
        return isinstance(other, PFuture) and self._fid == other._fid

    def __hash__(self) -> int:
        return self._fid

push/pfuture.py
lbai-lab/push
from multiprocessing import Queue
import torch
import torch.multiprocessing as mp
from typing import *

from push.lib.node_event_loop import NodeEventLoop
from push.lib.messages import *
from push.lib.waitable import Waitable

from push.pfuture import PFuture

    

def init_node_event_loop(mk_module: Callable,
                         args: List[any],
                         in_queues: Dict[int, Queue],
                         out_queues: Dict[int, Queue],
                         rank: int,
                         devices: List[int],
                         cache_size: int,
                         view_size: int) -> None:
    nel = NodeEventLoop(mk_module, args, in_queues[rank], out_queues[rank], rank, devices, cache_size, view_size)
    out_queues[rank].put(NodeEvtLoopInitMSG())
    nel._start_event_loop()


class PusH(Waitable):
    """PusH Distribution.

    1. Create a Push Distribution which approximates a distribution on nn's parameters via *particles*.
    2. Create arbitrary number of particles (pinit). Particles execute concurrently of other particles
    """    
    def __init__(self, mk_module: Callable, *args, cache_size=4, view_size=4) -> None:
        # Model
        self.mk_module = mk_module
        self.args = args

        # Process management
        try:
            mp.set_start_method("spawn")
        except:
            pass
        self._manager = mp.Manager()
        
        # Message queues for device event loops
        self._in_queues = {}             # device -> queue
        self._out_queues = {}            # device -> queue
        self._processes = {}             # device -> process
        self._particle_to_device = {}    # pid -> device
        self._particle_to_rank = {}      # pid -> rank
        
        # Device manager
        self.cache_size = cache_size
        self.view_size = view_size
        self.rank = 0
        self._init()

        # Tasks and results
        self._future_id = 0
        self._particle_to_futures = {-1: []}
        self._future_to_particle = {}
        self._results = {}

    def _init(self) -> None:
        # Create mailbox
        devices = []
        for device_id in range(torch.cuda.device_count()):
            devices += [device_id]
        
        self._in_queues[self.rank] = self._manager.Queue()
        self._out_queues[self.rank] = self._manager.Queue()

        # Start device event loops
        p = mp.Process(
            target=init_node_event_loop,
            args=(
                self.mk_module,
                self.args,
                self._in_queues,
                self._out_queues,
                self.rank,
                devices,
                self.cache_size,
                self.view_size,
            ))
        self._processes[self.rank] = p
        p.start()

        # Acknowledge that device event loops have been started
        msg = self._out_queues[self.rank].get()
        if not isinstance(msg, NodeEvtLoopInitMSG):
            raise ValueError(f"Fatal error ... inconsistent message state {msg}")

    # -----------------------------------------------------
    # Context management
    # -----------------------------------------------------

    def __enter__(self):
        return self

    def _cleanup(self) -> None:
        for device_id, proc in self._processes.items():
            self._in_queues[device_id].put(NodeEvtLoopCleanupMSG())
        for device_id, proc in self._processes.items():
            proc.join()

    def __exit__(self, exc_type, exc_value, traceback):
        self._cleanup()

    # -----------------------------------------------------
    # Helper functionality
    # -----------------------------------------------------

    def _create_future_id(self):
        fid = self._future_id
        self._future_id += 1
        return fid

    def _register_future(self, pid: int, fid: int) -> None:
        self._particle_to_futures[pid] += [fid]
        self._future_to_particle[fid] = pid

    def _pwait(self, fids: List[int]) -> Dict[int, any]:
        remaining = set(fids)
        acc = {}
        def loop():
            for fid in fids:
                if fid in self._results:
                    y = self._results.pop(fid)
                    acc[fid] = y
                    remaining.remove(fid)
                    pid = self._future_to_particle.pop(fid)
                    self._particle_to_futures[pid].remove(fid)
        loop()

        while len(remaining) > 0:
            for fid in fids:
                if fid in remaining:
                    pid = self._future_to_particle[fid]
                    rank = self._particle_to_rank[pid]
                    msg = self._out_queues[rank].get()

                    if isinstance(msg, Exception):
                        raise msg
                    elif isinstance(msg, ReceiveFuncAckPDMSG):
                        self._results[msg.pid_fid[1]] = None
                    loop()
                    break
        return acc

    def _wait(self, fid: int) -> any:
        raise NotImplementedError
    

    # -----------------------------------------------------
    # Particles API
    # -----------------------------------------------------

    def particle_ids(self) -> List[int]:
        """Returns all particles.

        Returns:
            List[int]: List of all particle identifiers visible to current particle.
        """
        return list(self._particle_to_rank.keys())

    def num_particles(self) -> int:
        """Returns number of particles.

        Returns:
            int: The number of particles.
        """        
        return len(self._particle_to_rank)

    def p_create(self, mk_optim: Callable, device=0, receive={}, state={}) -> int:
        """Create a particle

        Args:
            mk_optim (Callable): Optimizer for updating parameters. Can be None.
            device (int, optional): Device to put particle on. Defaults to 0.
            receive (dict, optional): 
               Dictionary containing messages that this particle will respond to.
               Defaults to {}.

        Returns:
            int: Particle identifier.
        """        
        # Create particle
        new_pid = len(self._particle_to_rank)
        self._particle_to_device[new_pid] = device
        self._particle_to_rank[new_pid] = self.rank
        self._particle_to_futures[new_pid] = []
        self._in_queues[self.rank].put(ReceiveParticleInitPDMSG(device, new_pid, mk_optim, receive, state))
        
        # Acknowledge
        msg = self._out_queues[self.rank].get()
        if not isinstance(msg, ReceiveParticleInitAckPDMSG):
            raise ValueError(f"Fatal error ... inconsistent message state")

        # Broadcast so we can discover other particles
        for pid, queue in self._in_queues.items():
            queue.put(NELBroadcastParticlesMSG(self._in_queues, self._out_queues, self._particle_to_device))
            msg = self._out_queues[pid].get()
            if not isinstance(msg, NELBroadcastParticlesAckMSG):
                raise ValueError(f"Fatal error ... inconsistent message state")
        
        return new_pid

    def p_parameters(self, pid: int, sync=True) -> Union[int, Iterable[torch.Tensor]]:
        """_summary_

        Args:
            pid (int): Identifier of particle to obtain parameters of.
            sync (bool, optional): Obtain parameters. Defaults to True.

        Returns:
            Union[int, Iterable[torch.Tensor]]: Parameters
        """        
        if pid not in self._particle_to_rank:
            raise ValueError(f"Particle {pid} does not exist")

        # Create task
        fid = self._create_future_id()
        self._register_future(pid, fid)

        # Initiate task
        rank = self._particle_to_rank[pid]
        self._in_queues[rank].put(ReceiveParametersPDMSG((pid, fid), pid))

        if sync:
            # Synchronize
            y = self._out_queues[rank].get()
            self._results[fid] = y
            while y.pid_fid[1] != fid:
                y = self._out_queues[rank].get()
                self._results[fid] = y
            return y.params
        else:
            # return fid
            return PFuture(self, -1, pid, fid)

    def p_launch(self, pid_to: int, msg: str, *args, sync=False) -> PFuture:
        """Launch a particle.

        Args:
            pid_to (int): Identifier of particle that is the main entry point.
            msg (str): Message associated with main function
            sync (bool, optional): Async. Defaults to False.

        Returns:
            PFuture: Calling wait will get the result of the computation.
        """        
        if pid_to not in self._particle_to_rank:
            raise ValueError(f"Particle {pid_to} does not exist")
        
        # Create task
        fid = self._create_future_id()
        self._register_future(pid_to, fid)

        # Send message
        rank = self._particle_to_rank[pid_to]
        self._in_queues[rank].put(ReceiveFuncPDMSG((pid_to, fid), pid_to, msg, args))

        if sync:
            # Synchronize
            y = self._out_queues[rank].get()
            self._results[fid] = y
            while y.pid_fid[1] != fid:
                y = self._out_queues[rank].get()
                self._results[fid] = y
            return None
        else:
            # return fid
            return PFuture(self, -1, pid_to, fid)

    def p_wait(self, futures: list[PFuture]) -> dict[int, any]:
        return self._pwait([future._fid for future in futures])

    # -----------------------------------------------------
    # Utility
    # -----------------------------------------------------

    def save(self, sync=True):
        pid = 0

        # Create task
        fid = self._create_future_id()
        self._register_future(pid, fid)

        # Initiate task
        rank = self._particle_to_rank[pid]
        self._in_queues[rank].put(NELSaveModel((pid, fid)))

        if sync:
            # Synchronize
            y = self._out_queues[rank].get()
            self._results[fid] = y
            while y.pid_fid[1] != fid:
                y = self._out_queues[rank].get()
                self._results[fid] = y
            return None
        else:
            # return fid
            return PFuture(self, -1, pid, fid)
        
push/push.py
lbai-lab/push
import pathlib
from pathlib import Path
from setuptools import setup, find_packages

# The directory containing this file
HERE = pathlib.Path(__file__).parent

# The text of the README file
README = (HERE / "README.md").read_text()


def get_install_requires() -> list[str]:
    """Returns requirements.txt parsed to a list"""
    fname = Path(__file__).parent / 'requirements.txt'
    targets = []
    if fname.exists():
        with open(fname, 'r') as f:
            targets = f.read().splitlines()
    print("targets =========")
    print(targets)
    return targets

# This call to setup() does all the work
setup(
    name="PusH",
    version="0.1.0",
    description="Concurrent Probabilistic Programming for Bayesian Deep Learning",
    long_description=README,
    long_description_content_type="text/markdown",
    url="https://github.com/lbai-lab/push/tree/main",
    author="The PusH Authors",
    author_email="dan.e.huang@gmail.com",
    license="Apache 2.0",
    classifiers=[
        "License :: OSI Approved :: Apache 2.0 License",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.10",
    ],
    packages=find_packages(),
    include_package_data=True,
    install_requires=get_install_requires(),
    extras_require={
    }
)

setup.py
lbai-lab/push
import argparse
from typing import *
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import push.bayes.ensemble
import push.bayes.stein_vgd
import push.bayes.swag

import time


torch.manual_seed(34)


"""[Note]

Quick Usage

1. `python test_basic.py -m ensemble`
2. `python test_basic.py -m stein_vgd`
3. `python test_basic.py -m mswag`

"""

# =============================================================================
# Simple Dataset + Neural Network
# =============================================================================

class RandDataset(Dataset):
    def __init__(self, batch_size, N, D):
        self.xs = torch.randn(batch_size*N, D)
        self.ys = torch.randn(batch_size*N, 1)

    def __len__(self):
        return len(self.xs)

    def __getitem__(self, idx):
        return self.xs[idx], self.ys[idx]


class MiniNN(nn.Module):
    def __init__(self, D):
        super(MiniNN, self).__init__()
        self.fc1 = nn.Linear(D, D)
        self.fc2 = nn.Linear(D, D)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.ReLU()(x)
        x = self.fc2(x)
        return x
    

class BiggerNN(nn.Module):
    def __init__(self, n, D):
        super(BiggerNN, self).__init__()
        self.minis = []
        self.n = n
        for i in range(0, n):
            self.minis += [MiniNN(D)]
            self.add_module("mini_layer"+str(i), self.minis[-1])
        self.fc = nn.Linear(D, 1)
            
    def forward(self, x):
        for i in range(0, self.n):
            x = self.minis[i](x)
        return self.fc(x)


# =============================================================================
# Basic Testing
# =============================================================================

if __name__ == "__main__":
    # Arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-L", "--layers", type=int, default=2,
                        help="size of neural network")
    parser.add_argument("-D", "--dimension", type=int, default=20,
                        help="size of neural network")
    parser.add_argument("-p", "--particles", type=int, default=2,
                        help="number of particles for SVGD or mswag")
    parser.add_argument("-b", "--batchsize", type=int, default=128,
                        help="number of batches")
    parser.add_argument("-e", "--epochs", type=int, default=10,
                        help="number of batches")
    parser.add_argument("-n", "--size", type=int, default=10,
                        help="size of dataset")
    parser.add_argument("-d", "--devices", type=int, default=1,
                        help="number of devices")
    parser.add_argument("-cs", "--cache_size", type=int, default=4,
                        help="size of the cache")
    parser.add_argument("-vs", "--view_size", type=int, default=4,
                        help="size of the cache")
    parser.add_argument("-m", "--method", type=str, default="ensemble",
                        choices=[
                            "ensemble",
                            "swag-central",
                            "swag",
                            "mswag",
                            "stein_vgd",
                        ],
                        help="bayesian inference method")
    args = parser.parse_args()

    # Dataset
    L = args.layers
    D = args.dimension
    N = args.size
    dataset = RandDataset(args.batchsize, N, D)
    dataloader = DataLoader(dataset, batch_size=args.batchsize, shuffle=True)
    
    # Dispatch
    if args.method == "ensemble":
        def mk_dataloader():
            return DataLoader(RandDataset(args.batchsize, args.size, args.dimension), batch_size=128, shuffle=True)

        epochs = 10
        num_ensembles = 3
        push.bayes.ensemble.train_deep_ensemble(
            dataloader, # mk_dataloader,
            torch.nn.MSELoss(),
            epochs,
            BiggerNN, L, D,
            num_devices=args.devices,
            num_ensembles=num_ensembles
        )

    elif args.method == "mswag":
        pretrain_epochs = 10
        swag_epochs = 5
        cache_size=4
        view_size=4
        start = time.perf_counter()
        push.bayes.swag.train_mswag(
            dataloader,
            torch.nn.MSELoss(),
            pretrain_epochs,
            swag_epochs,
            args.particles,
            cache_size,
            view_size, 
            BiggerNN, L, D
        )
        print("Time elapsed", time.perf_counter() - start)

    elif args.method == "stein_vgd":
        import numpy as np
        start = time.perf_counter()
        print("Size of model", sum([np.prod(p.size()) for p in BiggerNN(L, D).parameters()]))
        print("Size of dataset", len(dataloader.dataset))
        svgd_state = {
            "L": L,
            "D": D,
            "N": N,
        }
        push.bayes.stein_vgd.train_svgd(
            dataloader, torch.nn.MSELoss(),
            args.epochs, args.particles,
            BiggerNN, L, D,
            lengthscale=1.0, lr=1e-3, prior=None,
            cache_size=args.cache_size, view_size=args.view_size, num_devices=args.devices,
            svgd_entry=push.bayes.stein_vgd._svgd_leader, svgd_state=svgd_state
        )
        print("Time elapsed", time.perf_counter() - start)

    else:
        raise ValueError(f"Method {args.method} not supported ...")

test/test_basic.py
lbai-lab/push
import unittest
from typing import *
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import push.bayes.ensemble
import push.bayes.stein_vgd
import push.bayes.swag


# =============================================================================
# Simple Dataset + Neural Network
# =============================================================================

class RandDataset(Dataset):
    def __init__(self, D):
        self.xs = torch.randn(128*10, D)
        self.ys = torch.randn(128*10, 1)

    def __len__(self):
        return len(self.xs)

    def __getitem__(self, idx):
        return self.xs[idx], self.ys[idx]


class MiniNN(nn.Module):
    def __init__(self, D):
        super(MiniNN, self).__init__()
        self.fc1 = nn.Linear(D, D)
        self.fc2 = nn.Linear(D, D)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.ReLU()(x)
        x = self.fc2(x)
        return x
    

class BiggerNN(nn.Module):
    def __init__(self, n, D):
        super(BiggerNN, self).__init__()
        self.minis = []
        self.n = n
        for i in range(0, n):
            self.minis += [MiniNN(D)]
            self.add_module("mini_layer"+str(i), self.minis[-1])
        self.fc = nn.Linear(D, 1)
            
    def forward(self, x):
        for i in range(0, self.n):
            x = self.minis[i](x)
        return self.fc(x)


class TestEnd2End(unittest.TestCase):
    def test_ensemble(self):
        L = 10
        D = 20
        dataset = RandDataset(D)
        dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

        epochs = 10
        num_ensembles = 3
        push.bayes.ensemble.train_deep_ensemble(
            dataloader,
            torch.nn.MSELoss(),
            epochs,
            BiggerNN, L, D,
            num_ensembles=num_ensembles
        )

    def test_swag(self):
        L = 10
        D = 20
        dataset = RandDataset(D)
        dataloader = DataLoader(dataset, batch_size=128, shuffle=True)
        pretrain_epochs = 10
        swag_epochs = 5
        cache_size = 4
        view_size = 4
        num_models = 2
        push.bayes.swag.train_mswag(
            dataloader,
            torch.nn.MSELoss(),
            pretrain_epochs,
            swag_epochs,
            num_models,
            cache_size,
            view_size, 
            BiggerNN, L, D, num_devices=1
        )

    def test_stein_vgd(self):
        L = 10
        D = 20
        epochs = 3
        particles = 4
        devices = 1
        cache_size = 4
        view_size = 4
        dataset = RandDataset(D)
        dataloader = DataLoader(dataset, batch_size=128, shuffle=True)
        svgd_state = {
            "L": L,
            "D": D,
            "N": 10,
        }
        push.bayes.stein_vgd.train_svgd(
            dataloader, torch.nn.MSELoss(),
            epochs, particles,
            BiggerNN, L, D,
            lengthscale=1.0, lr=1e-3, prior=None,
            cache_size=cache_size, view_size=view_size, num_devices=devices,
            svgd_entry=push.bayes.stein_vgd._svgd_leader, svgd_state=svgd_state
        )

test/test_end2end.py
lbai-lab/push