wehub-resource-sync
309 lines
9.8 KiB
Python
309 lines
9.8 KiB
Python
import copy
|
|
from pathlib import Path
|
|
|
|
import click
|
|
|
|
import dgl
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
from dgl.data.utils import Subset
|
|
from logzero import logger
|
|
from modules.dimenet import DimeNet
|
|
from modules.dimenet_pp import DimeNetPP
|
|
from modules.initializers import GlorotOrthogonal
|
|
from qm9 import QM9
|
|
from ruamel.yaml import YAML
|
|
from sklearn.metrics import mean_absolute_error
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
def split_dataset(
|
|
dataset, num_train, num_valid, shuffle=False, random_state=None
|
|
):
|
|
"""Split dataset into training, validation and test set.
|
|
|
|
Parameters
|
|
----------
|
|
dataset
|
|
We assume that ``len(dataset)`` gives the number of datapoints and ``dataset[i]``
|
|
gives the ith datapoint.
|
|
num_train : int
|
|
Number of training datapoints.
|
|
num_valid : int
|
|
Number of validation datapoints.
|
|
shuffle : bool, optional
|
|
By default we perform a consecutive split of the dataset. If True,
|
|
we will first randomly shuffle the dataset.
|
|
random_state : None, int or array_like, optional
|
|
Random seed used to initialize the pseudo-random number generator.
|
|
This can be any integer between 0 and 2^32 - 1 inclusive, an array
|
|
(or other sequence) of such integers, or None (the default value).
|
|
If seed is None, then RandomState will try to read data from /dev/urandom
|
|
(or the Windows analogue) if available or seed from the clock otherwise.
|
|
|
|
Returns
|
|
-------
|
|
list of length 3
|
|
Subsets for training, validation and test.
|
|
"""
|
|
from itertools import accumulate
|
|
|
|
num_data = len(dataset)
|
|
assert num_train + num_valid < num_data
|
|
lengths = [num_train, num_valid, num_data - num_train - num_valid]
|
|
if shuffle:
|
|
indices = np.random.RandomState(seed=random_state).permutation(num_data)
|
|
else:
|
|
indices = np.arange(num_data)
|
|
return [
|
|
Subset(dataset, indices[offset - length : offset])
|
|
for offset, length in zip(accumulate(lengths), lengths)
|
|
]
|
|
|
|
|
|
@torch.no_grad()
|
|
def ema(ema_model, model, decay):
|
|
msd = model.state_dict()
|
|
for k, ema_v in ema_model.state_dict().items():
|
|
model_v = msd[k].detach()
|
|
ema_v.copy_(ema_v * decay + (1.0 - decay) * model_v)
|
|
|
|
|
|
def edge_init(edges):
|
|
R_src, R_dst = edges.src["R"], edges.dst["R"]
|
|
dist = torch.sqrt(F.relu(torch.sum((R_src - R_dst) ** 2, -1)))
|
|
# d: bond length, o: bond orientation
|
|
return {"d": dist, "o": R_src - R_dst}
|
|
|
|
|
|
def _collate_fn(batch):
|
|
graphs, line_graphs, labels = map(list, zip(*batch))
|
|
g, l_g = dgl.batch(graphs), dgl.batch(line_graphs)
|
|
labels = torch.tensor(labels, dtype=torch.float32)
|
|
return g, l_g, labels
|
|
|
|
|
|
def train(device, model, opt, loss_fn, train_loader):
|
|
model.train()
|
|
epoch_loss = 0
|
|
num_samples = 0
|
|
|
|
for g, l_g, labels in train_loader:
|
|
g = g.to(device)
|
|
l_g = l_g.to(device)
|
|
labels = labels.to(device)
|
|
logits = model(g, l_g)
|
|
loss = loss_fn(logits, labels.view([-1, 1]))
|
|
epoch_loss += loss.data.item() * len(labels)
|
|
num_samples += len(labels)
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
|
|
return epoch_loss / num_samples
|
|
|
|
|
|
@torch.no_grad()
|
|
def evaluate(device, model, valid_loader):
|
|
model.eval()
|
|
predictions_all, labels_all = [], []
|
|
|
|
for g, l_g, labels in valid_loader:
|
|
g = g.to(device)
|
|
l_g = l_g.to(device)
|
|
logits = model(g, l_g)
|
|
labels_all.extend(labels)
|
|
predictions_all.extend(
|
|
logits.view(
|
|
-1,
|
|
)
|
|
.cpu()
|
|
.numpy()
|
|
)
|
|
|
|
return np.array(predictions_all), np.array(labels_all)
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"-m",
|
|
"--model-cnf",
|
|
type=click.Path(exists=True),
|
|
help="Path of model config yaml.",
|
|
)
|
|
def main(model_cnf):
|
|
yaml = YAML(typ="safe")
|
|
model_cnf = yaml.load(Path(model_cnf))
|
|
model_name, model_params, train_params, pretrain_params = (
|
|
model_cnf["name"],
|
|
model_cnf["model"],
|
|
model_cnf["train"],
|
|
model_cnf["pretrain"],
|
|
)
|
|
logger.info(f"Model name: {model_name}")
|
|
logger.info(f"Model params: {model_params}")
|
|
logger.info(f"Train params: {train_params}")
|
|
|
|
if model_params["targets"] in ["mu", "homo", "lumo", "gap", "zpve"]:
|
|
model_params["output_init"] = nn.init.zeros_
|
|
else:
|
|
# 'GlorotOrthogonal' for alpha, R2, U0, U, H, G, and Cv
|
|
model_params["output_init"] = GlorotOrthogonal
|
|
|
|
logger.info("Loading Data Set")
|
|
dataset = QM9(label_keys=model_params["targets"], edge_funcs=[edge_init])
|
|
|
|
# data split
|
|
train_data, valid_data, test_data = split_dataset(
|
|
dataset,
|
|
num_train=train_params["num_train"],
|
|
num_valid=train_params["num_valid"],
|
|
shuffle=True,
|
|
random_state=train_params["data_seed"],
|
|
)
|
|
logger.info(f"Size of Training Set: {len(train_data)}")
|
|
logger.info(f"Size of Validation Set: {len(valid_data)}")
|
|
logger.info(f"Size of Test Set: {len(test_data)}")
|
|
|
|
# data loader
|
|
train_loader = DataLoader(
|
|
train_data,
|
|
batch_size=train_params["batch_size"],
|
|
shuffle=True,
|
|
collate_fn=_collate_fn,
|
|
num_workers=train_params["num_workers"],
|
|
)
|
|
|
|
valid_loader = DataLoader(
|
|
valid_data,
|
|
batch_size=train_params["batch_size"],
|
|
shuffle=False,
|
|
collate_fn=_collate_fn,
|
|
num_workers=train_params["num_workers"],
|
|
)
|
|
|
|
test_loader = DataLoader(
|
|
test_data,
|
|
batch_size=train_params["batch_size"],
|
|
shuffle=False,
|
|
collate_fn=_collate_fn,
|
|
num_workers=train_params["num_workers"],
|
|
)
|
|
|
|
# check cuda
|
|
gpu = train_params["gpu"]
|
|
device = f"cuda:{gpu}" if gpu >= 0 and torch.cuda.is_available() else "cpu"
|
|
|
|
# model initialization
|
|
logger.info("Loading Model")
|
|
if model_name == "dimenet":
|
|
model = DimeNet(
|
|
emb_size=model_params["emb_size"],
|
|
num_blocks=model_params["num_blocks"],
|
|
num_bilinear=model_params["num_bilinear"],
|
|
num_spherical=model_params["num_spherical"],
|
|
num_radial=model_params["num_radial"],
|
|
cutoff=model_params["cutoff"],
|
|
envelope_exponent=model_params["envelope_exponent"],
|
|
num_before_skip=model_params["num_before_skip"],
|
|
num_after_skip=model_params["num_after_skip"],
|
|
num_dense_output=model_params["num_dense_output"],
|
|
num_targets=len(model_params["targets"]),
|
|
output_init=model_params["output_init"],
|
|
).to(device)
|
|
elif model_name == "dimenet++":
|
|
model = DimeNetPP(
|
|
emb_size=model_params["emb_size"],
|
|
out_emb_size=model_params["out_emb_size"],
|
|
int_emb_size=model_params["int_emb_size"],
|
|
basis_emb_size=model_params["basis_emb_size"],
|
|
num_blocks=model_params["num_blocks"],
|
|
num_spherical=model_params["num_spherical"],
|
|
num_radial=model_params["num_radial"],
|
|
cutoff=model_params["cutoff"],
|
|
envelope_exponent=model_params["envelope_exponent"],
|
|
num_before_skip=model_params["num_before_skip"],
|
|
num_after_skip=model_params["num_after_skip"],
|
|
num_dense_output=model_params["num_dense_output"],
|
|
num_targets=len(model_params["targets"]),
|
|
extensive=model_params["extensive"],
|
|
output_init=model_params["output_init"],
|
|
).to(device)
|
|
else:
|
|
raise ValueError(f"Invalid Model Name {model_name}")
|
|
|
|
if pretrain_params["flag"]:
|
|
torch_path = pretrain_params["path"]
|
|
target = model_params["targets"][0]
|
|
model.load_state_dict(
|
|
torch.load(f"{torch_path}/{target}.pt", weights_only=False)
|
|
)
|
|
|
|
logger.info("Testing with Pretrained model")
|
|
predictions, labels = evaluate(device, model, test_loader)
|
|
test_mae = mean_absolute_error(labels, predictions)
|
|
logger.info(f"Test MAE {test_mae:.4f}")
|
|
|
|
return
|
|
# define loss function and optimization
|
|
loss_fn = nn.L1Loss()
|
|
opt = optim.Adam(
|
|
model.parameters(),
|
|
lr=train_params["lr"],
|
|
weight_decay=train_params["weight_decay"],
|
|
amsgrad=True,
|
|
)
|
|
scheduler = optim.lr_scheduler.StepLR(
|
|
opt, train_params["step_size"], gamma=train_params["gamma"]
|
|
)
|
|
|
|
# model training
|
|
best_mae = 1e9
|
|
no_improvement = 0
|
|
|
|
# EMA for valid and test
|
|
logger.info("EMA Init")
|
|
ema_model = copy.deepcopy(model)
|
|
for p in ema_model.parameters():
|
|
p.requires_grad_(False)
|
|
|
|
best_model = copy.deepcopy(ema_model)
|
|
|
|
logger.info("Training")
|
|
for i in range(train_params["epochs"]):
|
|
train_loss = train(device, model, opt, loss_fn, train_loader)
|
|
ema(ema_model, model, train_params["ema_decay"])
|
|
if i % train_params["interval"] == 0:
|
|
predictions, labels = evaluate(device, ema_model, valid_loader)
|
|
|
|
valid_mae = mean_absolute_error(labels, predictions)
|
|
logger.info(
|
|
f"Epoch {i} | Train Loss {train_loss:.4f} | Val MAE {valid_mae:.4f}"
|
|
)
|
|
|
|
if valid_mae > best_mae:
|
|
no_improvement += 1
|
|
if no_improvement == train_params["early_stopping"]:
|
|
logger.info("Early stop.")
|
|
break
|
|
else:
|
|
no_improvement = 0
|
|
best_mae = valid_mae
|
|
best_model = copy.deepcopy(ema_model)
|
|
else:
|
|
logger.info(f"Epoch {i} | Train Loss {train_loss:.4f}")
|
|
|
|
scheduler.step()
|
|
|
|
logger.info("Testing")
|
|
predictions, labels = evaluate(device, best_model, test_loader)
|
|
test_mae = mean_absolute_error(labels, predictions)
|
|
logger.info("Test MAE {:.4f}".format(test_mae))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|