Convert v1 to v2

# Install chemprop from GitHub if running in Google Colab
import os

if os.getenv("COLAB_RELEASE_TAG"):
    try:
        import chemprop
    except ImportError:
        !git clone https://github.com/chemprop/chemprop.git
        %cd chemprop
        !pip install .
        %cd examples

Import packages

import torch
from pprint import pprint
from pathlib import Path

from chemprop.utils.v1_to_v2 import convert_model_dict_v1_to_v2
from chemprop.models.model import MPNN

Change model paths here

chemprop_dir = Path.cwd().parent
model_v1_input_path =  chemprop_dir / "tests/data/example_model_v1_regression_mol.pt" # path to v1 model .pt file
model_v2_output_path = Path.cwd() / "converted_model.ckpt" # path to save the converted model .ckpt file

Load v1 model .pt file

model_v1_dict = torch.load(model_v1_input_path, weights_only=False)

# Here are all the keys that is stored in v1 model
pprint(list(model_v1_dict.keys()))

['args',
 'state_dict',
 'data_scaler',
 'features_scaler',
 'atom_descriptor_scaler',
 'bond_descriptor_scaler',
 'atom_bond_scaler']

# Here are the input arguments that is stored in v1 model
pprint(model_v1_dict['args'].__dict__)

{'activation': 'ReLU',
 'adding_bond_types': True,
 'adding_h': False,
 'aggregation': 'mean',
 'aggregation_norm': 100,
 'atom_constraints': [],
 'atom_descriptor_scaling': True,
 'atom_descriptors': None,
 'atom_descriptors_path': None,
 'atom_descriptors_size': 0,
 'atom_features_size': 0,
 'atom_messages': False,
 'atom_targets': [],
 'batch_size': 50,
 'bias': False,
 'bias_solvent': False,
 'bond_constraints': [],
 'bond_descriptor_scaling': True,
 'bond_descriptors': None,
 'bond_descriptors_path': None,
 'bond_descriptors_size': 0,
 'bond_features_size': 0,
 'bond_targets': [],
 'cache_cutoff': 10000,
 'checkpoint_dir': None,
 'checkpoint_frzn': None,
 'checkpoint_path': None,
 'checkpoint_paths': None,
 'class_balance': False,
 'config_path': None,
 'constraints_path': None,
 'crossval_index_dir': None,
 'crossval_index_file': None,
 'crossval_index_sets': None,
 'cuda': False,
 'data_path': '/Users/hwpang/Software/chemprop/tests/data/regression.csv',
 'data_weights_path': None,
 'dataset_type': 'regression',
 'depth': 3,
 'depth_solvent': 3,
 'device': device(type='cpu'),
 'dropout': 0.0,
 'empty_cache': False,
 'ensemble_size': 1,
 'epochs': 1,
 'evidential_regularization': 0,
 'explicit_h': False,
 'extra_metrics': [],
 'features_generator': None,
 'features_only': False,
 'features_path': None,
 'features_scaling': True,
 'features_size': None,
 'ffn_hidden_size': 300,
 'ffn_num_layers': 2,
 'final_lr': 0.0001,
 'folds_file': None,
 'freeze_first_only': False,
 'frzn_ffn_layers': 0,
 'gpu': None,
 'grad_clip': None,
 'hidden_size': 300,
 'hidden_size_solvent': 300,
 'ignore_columns': None,
 'init_lr': 0.0001,
 'is_atom_bond_targets': False,
 'keeping_atom_map': False,
 'log_frequency': 10,
 'loss_function': 'mse',
 'max_data_size': None,
 'max_lr': 0.001,
 'metric': 'rmse',
 'metrics': ['rmse'],
 'minimize_score': True,
 'mpn_shared': False,
 'multiclass_num_classes': 3,
 'no_adding_bond_types': False,
 'no_atom_descriptor_scaling': False,
 'no_bond_descriptor_scaling': False,
 'no_cache_mol': False,
 'no_cuda': False,
 'no_features_scaling': False,
 'no_shared_atom_bond_ffn': False,
 'num_folds': 1,
 'num_lrs': 1,
 'num_tasks': 1,
 'num_workers': 8,
 'number_of_molecules': 1,
 'overwrite_default_atom_features': False,
 'overwrite_default_bond_features': False,
 'phase_features_path': None,
 'pytorch_seed': 0,
 'quiet': False,
 'reaction': False,
 'reaction_mode': 'reac_diff',
 'reaction_solvent': False,
 'resume_experiment': False,
 'save_dir': '/Users/hwpang/Software/test_chemprop_v1_to_v2/fold_0',
 'save_preds': False,
 'save_smiles_splits': True,
 'seed': 0,
 'separate_test_atom_descriptors_path': None,
 'separate_test_bond_descriptors_path': None,
 'separate_test_constraints_path': None,
 'separate_test_features_path': None,
 'separate_test_path': None,
 'separate_test_phase_features_path': None,
 'separate_val_atom_descriptors_path': None,
 'separate_val_bond_descriptors_path': None,
 'separate_val_constraints_path': None,
 'separate_val_features_path': None,
 'separate_val_path': None,
 'separate_val_phase_features_path': None,
 'shared_atom_bond_ffn': True,
 'show_individual_scores': False,
 'smiles_columns': ['smiles'],
 'spectra_activation': 'exp',
 'spectra_phase_mask': None,
 'spectra_phase_mask_path': None,
 'spectra_target_floor': 1e-08,
 'split_key_molecule': 0,
 'split_sizes': [0.8, 0.1, 0.1],
 'split_type': 'random',
 'target_columns': None,
 'target_weights': None,
 'task_names': ['logSolubility'],
 'test': False,
 'test_fold_index': None,
 'train_data_size': 400,
 'undirected': False,
 'use_input_features': False,
 'val_fold_index': None,
 'warmup_epochs': 2.0,
 'weights_ffn_num_layers': 2}

# Here are the state_dict that is stored in v1 model
pprint(list(model_v1_dict['state_dict'].keys()))

['encoder.encoder.0.cached_zero_vector',
 'encoder.encoder.0.W_i.weight',
 'encoder.encoder.0.W_h.weight',
 'encoder.encoder.0.W_o.weight',
 'encoder.encoder.0.W_o.bias',
 'readout.1.weight',
 'readout.1.bias',
 'readout.4.weight',
 'readout.4.bias']

Convert loaded v1 model dictionary into v2 model dictionary

model_v2_dict = convert_model_dict_v1_to_v2(model_v1_dict)

# Here are all the keys in the converted model
pprint(list(model_v2_dict.keys()))

['epoch',
 'global_step',
 'pytorch-lightning_version',
 'state_dict',
 'loops',
 'callbacks',
 'optimizer_states',
 'lr_schedulers',
 'hparams_name',
 'hyper_parameters']

# Here are all the keys in the converted state_dict
pprint(list(model_v2_dict['state_dict'].keys()))

['message_passing.W_i.weight',
 'message_passing.W_h.weight',
 'message_passing.W_o.weight',
 'message_passing.W_o.bias',
 'predictor.ffn.0.0.weight',
 'predictor.ffn.0.0.bias',
 'predictor.ffn.1.2.weight',
 'predictor.ffn.1.2.bias',
 'predictor.output_transform.mean',
 'predictor.output_transform.scale',
 'predictor.criterion.task_weights']

# Here are all the keys in the converted hyper_parameters
pprint(list(model_v2_dict['hyper_parameters'].keys()))

['batch_norm',
 'metrics',
 'warmup_epochs',
 'init_lr',
 'max_lr',
 'final_lr',
 'message_passing',
 'agg',
 'predictor']

Save

torch.save(model_v2_dict, model_v2_output_path)

Load converted model

mpnn = MPNN.load_from_checkpoint(model_v2_output_path)

# now visually check the converted model is what is expected
mpnn

MPNN(
  (message_passing): BondMessagePassing(
    (W_i): Linear(in_features=147, out_features=300, bias=False)
    (W_h): Linear(in_features=300, out_features=300, bias=False)
    (W_o): Linear(in_features=433, out_features=300, bias=True)
    (dropout): Dropout(p=0.0, inplace=False)
    (tau): ReLU()
    (V_d_transform): Identity()
    (graph_transform): Identity()
  )
  (agg): MeanAggregation()
  (bn): Identity()
  (predictor): RegressionFFN(
    (ffn): MLP(
      (0): Sequential(
        (0): Linear(in_features=300, out_features=300, bias=True)
      )
      (1): Sequential(
        (0): ReLU()
        (1): Dropout(p=0.0, inplace=False)
        (2): Linear(in_features=300, out_features=1, bias=True)
      )
    )
    (criterion): MSE(task_weights=[[1.0]])
    (output_transform): UnscaleTransform()
  )
  (X_d_transform): Identity()
  (metrics): ModuleList(
    (0): RMSE(task_weights=[[1.0]])
    (1): MSE(task_weights=[[1.0]])
  )
)