Fail to load state dict when running sample.py

[Yang Fan]

Hi,
I am trying to train the model on a new dataset, so I run train.py and sample.py with new configuration files.

My training and sampling commands:

python train.py my_train_main_model.yml
python train.py my_train_frontier_model.yml
python sample.py my_sample.yml -i 0

my_train_main_model.yml, my_train_frontier_model.yml only change dataset path, and my_sample.yml change checkpoint and dataset paths.

However, I got following error message:

[2022-02-22 07:56:33,069::sample::INFO] Namespace(config='configs/fy_sample.yml', data_id=0, device='cuda', outdir='./outputs')
[2022-02-22 07:56:33,069::sample::INFO] {'dataset': {'name': 'pl', 'path': '/home/t-yafan/workspace/data/Tgt2Drug', 'split': '/home/t-yafan/workspace/data/Tgt2Drug/Binary/v4-2/src/split_by_name.pt'}, 'model': {'main': {'checkpoint': './logs/fy_train_main_model_2022_02_21__07_57_11/checkpoints/100000.pt'}, 'frontier': {'checkpoint': './logs/fy_train_frontier_model_2022_02_21__14_30_45/checkpoints/80000.pt'}}, 'sample': {'seed': 2020, 'num_samples': 100, 'beam_size': 300, 'logp_thres': -inf, 'num_retry': 5, 'max_steps': 50}}
[2022-02-22 07:56:33,070::sample::INFO] Loading data...
[2022-02-22 07:56:33,632::sample::INFO] Loading main model...
[2022-02-22 07:56:36,019::sample::INFO] Loading frontier model...
Traceback (most recent call last):
  File "sample.py", line 194, in <module>
    ftnet.load_state_dict(ckpt_ft['model'])
  File "/anaconda/envs/SBDD-3D/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1482, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for FrontierNetwork:
        Missing key(s) in state_dict: "frontier_pred.layers.0.weight", "frontier_pred.layers.0.bias", "frontier_pred.layers.1.weight", "frontier_pred.layers.1.bias", "frontier_pred.layers.2.weight", "frontier_pred.layers.2.bias", "frontier_pred.layers.3.weight", "frontier_pred.layers.3.bias". 
        Unexpected key(s) in state_dict: "field.lin1.weight", "field.lin2.weight", "field.lin2.bias", "field.nn.0.weight", "field.nn.0.bias", "field.nn.2.weight", "field.nn.2.bias", "field.classifier.0.weight", "field.classifier.0.bias", "field.classifier.2.weight", "field.classifier.2.bias", "field.property_pred.0.weight", "field.property_pred.0.bias", "field.property_pred.2.weight", "field.property_pred.2.bias", "field.distance_expansion.offset".

Which step goes wrong? How can I sample on trained models?

Thanks!

[Shitong Luo]

It seems that you are loading main-model weights to the frontier model, according to the repored missing keys and unexpected keys:

        Missing key(s) in state_dict: "frontier_pred.layers.0.weight", "frontier_pred.layers.0.bias", "frontier_pred.layers.1.weight", "frontier_pred.layers.1.bias", "frontier_pred.layers.2.weight", "frontier_pred.layers.2.bias", "frontier_pred.layers.3.weight", "frontier_pred.layers.3.bias". 
        Unexpected key(s) in state_dict: "field.lin1.weight", "field.lin2.weight", "field.lin2.bias", "field.nn.0.weight", "field.nn.0.bias", "field.nn.2.weight", "field.nn.2.bias", "field.classifier.0.weight", "field.classifier.0.bias", "field.classifier.2.weight", "field.classifier.2.bias", "field.property_pred.0.weight", "field.property_pred.0.bias", "field.property_pred.2.weight", "field.property_pred.2.bias", "field.distance_expansion.offset".

[Shitong Luo]

It should be the following command to train a frontier model:

python train_frontier.py my_train_frontier_model.yml

[Yang Fan]

Very thanks!