No module named 'torch_geometric.data.storage'

[hkim716]

When I tried to load my own graph data, there was an error about torch_geometric.data.storage.
I think patch_graph_all_v1.pt was generated from another PyG environment that has different package versions.

----> 3 gg_trn = torch.load(os.path.join(path_data,'patch_graph_all_v1.pt'))
      4 trn_ldr = DataLoader(gg_trn,batch_size=1000,shuffle=True)

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch/serialization.py in load(f, map_location, pickle_module, **pickle_load_args)
    605                     opened_file.seek(orig_position)
    606                     return torch.jit.load(opened_file)
--> 607                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
    608         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
    609 

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch/serialization.py in _load(zip_file, map_location, pickle_module, pickle_file, **pickle_load_args)
    880     unpickler = UnpicklerWrapper(data_file, **pickle_load_args)
    881     unpickler.persistent_load = persistent_load
--> 882     result = unpickler.load()
    883 
    884     torch._utils._validate_loaded_sparse_tensors()

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch/serialization.py in find_class(self, mod_name, name)
    873         def find_class(self, mod_name, name):
    874             mod_name = load_module_mapping.get(mod_name, mod_name)
--> 875             return super().find_class(mod_name, name)
    876 
    877     # Load the data (which may in turn use `persistent_load` to load tensors)

ModuleNotFoundError: No module named 'torch_geometric.data.storage'

I don't know how to solve this error. I think current torch_geometric is dependent on certain version of torch_spline_conv and so on, which means I cannot easily upgrade the version of current torch_geometric.

My environment has these versions below.

Current Python Version- 3.8.11
torch.__version___1.9.1
cuda.__version___10.2
torch_geometric.__version___1.7.2
torch_scatter.__version___2.0.8
torch_spline_conv.__version___1.2.1
torch_sparse.__version___0.6.12
torch_cluster.__version___1.5.9

[rusty1s]

Yes, it looks like you need torch-geometric>=2.0.0 to load this data in. This should be achievable by just running:

pip install --upgrade torch-geometric

No need to do any changes to the other packages.

[hkim716]

Thanks Matt, It worked!

I have another question though which is about how we save a trained model.
I used to use torch.save(model,os.path.join(path_model,'saved_model.pt')) to save a model, but I had some error when I load the saved model.
People recommended me to use torch.save(model.state_dict(), 'saved_model.pt') so I can load like this below.
model = GAE()
model.load_state_dict(torch.load(os.path.join(path_model,'saved_model.pt'))).

Why is it better way compared to the first one?

[rusty1s]

Yes, the state_dict approach is indeed recommended (see here):

This save/load process uses the most intuitive syntax and involves the least amount of code. Saving a model in this way will save the entire module using Python’s pickle module. The disadvantage of this approach is that the serialized data is bound to the specific classes and the exact directory structure used when the model is saved. The reason for this is because pickle does not save the model class itself. Rather, it saves a path to the file containing the class, which is used during load time. Because of this, your code can break in various ways when used in other projects or after refactors.