diff --git a/cgcnn2/cli/cgcnn_ft.py b/cgcnn2/cli/cgcnn_ft.py index 0dce4e9..1af81d1 100644 --- a/cgcnn2/cli/cgcnn_ft.py +++ b/cgcnn2/cli/cgcnn_ft.py @@ -12,7 +12,6 @@ from cgcnn2.data import CIFData, collate_pool, full_set_split from cgcnn2.model import CrystalGraphConvNet from cgcnn2.utils import ( - Normalizer, cgcnn_test, get_lr, print_checkpoint_info, @@ -293,7 +292,7 @@ def main(): sys.exit(1) # Load checkpoint onto device checkpoint = torch.load( - args.model_path, map_location=args.device, weights_only=False + args.model_path, map_location=args.device, weights_only=True ) model_args = argparse.Namespace(**checkpoint["args"]) @@ -321,9 +320,6 @@ def main(): model.to(args.device) model.eval() - normalizer = Normalizer(torch.zeros(3)) - normalizer.load_state_dict(checkpoint["normalizer"]) - print_checkpoint_info(checkpoint, args.model_path) # Initialize DataLoader @@ -522,7 +518,6 @@ def main(): savepoint = { "epoch": epoch + 1, "state_dict": model.state_dict(), - "normalizer": normalizer.state_dict(), "best_mse_error": avg_valid_loss, "args": vars(model_args), } @@ -547,7 +542,7 @@ def main(): # TEST WITH BEST MODEL # -------------------- checkpoint = torch.load( - os.path.join(output_folder, "best_model.ckpt"), weights_only=False + os.path.join(output_folder, "best_model.ckpt"), weights_only=True ) model.load_state_dict(checkpoint["state_dict"]) diff --git a/cgcnn2/cli/cgcnn_pr.py b/cgcnn2/cli/cgcnn_pr.py index 6d8b7a0..69b9151 100644 --- a/cgcnn2/cli/cgcnn_pr.py +++ b/cgcnn2/cli/cgcnn_pr.py @@ -124,7 +124,7 @@ def main(): # Load checkpoint onto device checkpoint = torch.load( - args.model_path, map_location=args.device, weights_only=False + args.model_path, map_location=args.device, weights_only=True ) model_args = argparse.Namespace(**checkpoint["args"]) diff --git a/cgcnn2/cli/cgcnn_tr.py b/cgcnn2/cli/cgcnn_tr.py index c815b2c..2261e11 100644 --- a/cgcnn2/cli/cgcnn_tr.py +++ b/cgcnn2/cli/cgcnn_tr.py @@ -2,7 +2,6 @@ import logging import os from pprint import pformat -from random import sample import sys import torch @@ -11,7 +10,7 @@ from cgcnn2.data import CIFData, collate_pool, full_set_split from cgcnn2.model import CrystalGraphConvNet -from cgcnn2.utils import Normalizer, cgcnn_test, get_lr, seed_everything, setup_logging +from cgcnn2.utils import cgcnn_test, get_lr, seed_everything, setup_logging def parse_arguments(args=None): @@ -275,26 +274,6 @@ def main(): ) sys.exit(1) - full_dataset = [*train_dataset, *valid_dataset, *test_dataset] - - # Normalizer setup - # For classification we use a dummy normalizer, otherwise compute mean/std from data - if args.task == "classification": - normalizer = Normalizer(torch.zeros(2)) - normalizer.load_state_dict({"mean": 0.0, "std": 1.0}) - else: - if len(full_dataset) < 500: - logging.warning( - "Dataset has fewer than 500 data points; results may have higher variance." - ) - sample_data_list = [full_dataset[i] for i in range(len(full_dataset))] - else: - sample_indices = sample(range(len(full_dataset)), 500) - sample_data_list = [full_dataset[i] for i in sample_indices] - - _, sample_target, _ = collate_pool(sample_data_list) - normalizer = Normalizer(sample_target) - # Build model # 1) gather input dimensions from first sample structures, _, _ = train_dataset[0] @@ -446,7 +425,6 @@ def main(): savepoint = { "epoch": epoch + 1, "state_dict": model.state_dict(), - "normalizer": normalizer.state_dict(), "best_mse_error": avg_valid_loss, "args": vars(args), } @@ -470,10 +448,9 @@ def main(): # TEST WITH BEST MODEL # -------------------- best_checkpoint = torch.load( - os.path.join(output_folder, "best_model.ckpt"), weights_only=False + os.path.join(output_folder, "best_model.ckpt"), weights_only=True ) model.load_state_dict(best_checkpoint["state_dict"]) - normalizer.load_state_dict(best_checkpoint["normalizer"]) cgcnn_test( model, diff --git a/cgcnn2/utils.py b/cgcnn2/utils.py index 6f40662..9b0fee7 100644 --- a/cgcnn2/utils.py +++ b/cgcnn2/utils.py @@ -708,8 +708,8 @@ def cgcnn_pred( checkpoint = torch.load( model_path, - map_location=lambda storage, loc: storage if not cuda else None, - weights_only=False, + map_location="cuda" if cuda else "cpu", + weights_only=True, ) structures, _, _ = total_dataset[0] orig_atom_fea_len = structures[0].shape[-1] @@ -726,8 +726,6 @@ def cgcnn_pred( if cuda: model.cuda() - normalizer = Normalizer(torch.zeros(3)) - normalizer.load_state_dict(checkpoint["normalizer"]) model.load_state_dict(checkpoint["state_dict"]) if verbose >= 100: diff --git a/docs/3_usage.md b/docs/3_usage.md index 9429c06..bb63657 100644 --- a/docs/3_usage.md +++ b/docs/3_usage.md @@ -72,7 +72,7 @@ Besides, we need some information about the pre-trained model architecture, whic import torch import argparse -checkpoint = torch.load(args.model_path, map_location=args.device, weights_only=False) +checkpoint = torch.load(args.model_path, map_location=args.device, weights_only=True) model_args = argparse.Namespace(**checkpoint["args"]) atom_fea_len = model_args.atom_fea_len n_conv = model_args.n_conv @@ -80,7 +80,7 @@ h_fea_len = model_args.h_fea_len n_h = model_args.n_h ``` -where `atom_fea_len`, `n_conv`, `h_fea_len`, and `n_h` are the dimensions of the atom features, the number of convolutional layers, the dimension of the hidden features, and the number of hidden layers, respectively. Since PyTorch 2.6, `torch.load` defaults to `weights_only=True`, so `weights_only=False` is required to load full CGCNN checkpoints; only load checkpoint files from sources you trust. Now, we can initialize the model by: +where `atom_fea_len`, `n_conv`, `h_fea_len`, and `n_h` are the dimensions of the atom features, the number of convolutional layers, the dimension of the hidden features, and the number of hidden layers, respectively. CGCNN checkpoints contain only tensors and plain Python values, so they load with `weights_only=True` (the secure default since PyTorch 2.6), which prevents checkpoint files from executing arbitrary code when loaded. Now, we can initialize the model by: ```python model = CrystalGraphConvNet(