Jie-Cao/AtomNet main
Module built on the numeric Python stackImports the NumPy array library
SHA-256c0ff166c9b69ee0c391812e03d1a3af532e7e2c53c1990cdad16a8c40c05ba93
MaleculeH₂(Db₂F)
Evidence
⋯4 lines
5import time
6import os
7import numpy as np
8import os.path as osp
9from tqdm import tqdm
10from torch_geometric.graphgym.config import cfg
11from torch.optim.lr_scheduler …
77:15… current best stats on eval epoch.
78 best_epoch = int(np.array([vp['MAE'] for vp in perf[1]]).argmin())
79 best_train = f"train_MAE: {perf[0][best_epoch]['MAE']:.4f}"
80 best_val = …
86:23… , step=cur_epoch)
87
88 run.summary["full_epoch_time_avg"] = np.mean(full_epoch_times)
89 run.summary["full_epoch_time_sum"] = np.sum(full_epoch_times)
90
91 if cfg.useSWA and cur_epoch + cfg.swa_epochs >= cfg.optim.max_epoch: # SWA
92 swa_model.up …
⋯4 lines
100
101 os.makedirs(ckpt_dir, exist_ok=True) # ckpt_dir = "results/jarvis_dft_3D_formation_energy_peratom/seed/ckpt/"
102 ckpt_path = osp.join(ckpt_dir, 'best.ckpt')
103
104 torch.save(ckpt, ckpt_path)
105
106 logging.info(f"Best checkpoint saved at {ckpt_path}" …
121:24… model.state_dict(), osp.join(ckpt_dir, 'swa.ckpt'))
122
123 ckpt = torch.load(ckpt_path)
124 model.load_state_dict(ckpt["model_state"])
125
126 eval_epoch(loggers[-1], loaders[-1], model, test_metrics=True)
127