Loading a model saved on PyG 2.3 using PyG 2.4

[ver228]

Hello,
I am trying to load the state dictionary on a model saved using PyG 2.3 after updating to PyG 2.4. The model is created using HeteroConv with lazy initialisation. The problem is that when i try to use torch.load I get a bunch of this type of errors:

 RuntimeError: Error(s) in loading state_dict for EdgePredictor:
E               Missing key(s) in state_dict: "encoder.layers.0.convs.<gene___evidence___disease>.lin_l.weight", "encoder.layers.0.convs.<gene___evidence___disease>.lin_l.bias", 
...
E               Unexpected key(s) in state_dict: "encoder.layers.0.convs.gene__evidence__disease.lin_l.weight", "encoder.layers.0.convs.gene__evidence__disease.lin_l.bias",
...

It just look to me that the way an edge type is saved in the state dict changed from "node1__edge__node2" to "<node1___edge___node2>". I did a fix to correct the keys to the new format but after the errors becomes something like:

While copying the parameter named "encoder.layers.0.convs.<gene___evidence___disease>.lin_l.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 15]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (15) at non-singleton dimension 1',).

Any suggestion on how to fix that or what is the reason would be appreciated.

[akihironitta]

@ver228 Would you mind sharing your code and env details to reproduce the behaviour? I'm not sure yet but fb1d855 looks suspicious.

[ver228]

@akihironitta Sorry the delay I had a few deadlines and until now i had time to create a minimal example. The problem is a bit worse that i thought in my opinion. It seems that in torch_geometric==2.4 one cannot load a state dictionary if the model uses lazy initialisation. My minimal example is:

import torch
import torch_geometric
from torch import nn
from torch_geometric.nn import HeteroConv, SAGEConv
from torch_geometric.data import HeteroData

class Encoder(nn.Module):
    def __init__(
        self,
        channels,
        edges
    ) -> None:
        super().__init__()
        self.layer1 = HeteroConv({k: SAGEConv((-1, -1), channels, aggr= "max") for k in edges}, aggr="mean")
        self.layer2 = HeteroConv({k: SAGEConv((-1, -1), channels, aggr= "max") for k in edges}, aggr="mean")

    def forward(self, x, edge_index):
        x = self.layer1(x, edge_index)
        x = {key: c.relu() for key, c in x.items()}
        x_out = self.layer2(x, edge_index)
        return x_out

def _make_fake_data():
    
    g = torch.Generator().manual_seed(42)
    def _make_edge_index(n_node1, n_node2, n_edges) -> torch.Tensor:
        out = torch.stack(
            (
                torch.randint(0, n_node1 - 1, (n_edges,), generator=g),
                torch.randint(0, n_node2 - 1, (n_edges,), generator=g),
            )
        )
        out: torch.Tensor = torch.unique(out, dim=1)
        return out
    
    N_GENES = 50
    N_DISEASES = 200
    data = HeteroData(
        {
            "disease": {"x": torch.zeros(N_DISEASES, 5)},
            "gene": {"x": torch.zeros(N_GENES, 5)},
            ("gene", "to", "disease"):  {"edge_index": _make_edge_index(N_GENES, N_DISEASES, 100)},
            ("disease", "to", "gene"):  {"edge_index": _make_edge_index(N_DISEASES, N_GENES, 100)},
        }
    )
    return data


if __name__ == "__main__":
    save_postfix = ''.join(torch_geometric.__version__.split("."))

    data = _make_fake_data()

    _, edges = data.metadata()

    net1 = Encoder(32, edges)
    out = net1(data.x_dict, data.edge_index_dict)

    torch.save(net1.state_dict(), f"test_{save_postfix}.pt")
    
    state = torch.load(f"test_{save_postfix}.pt")
    
    net2 = Encoder(32, edges)
    net2.load_state_dict(state)


    

I tried to environments one for torch_geometric==2.3 and another with torch_geometric==2.4. They are basically identical except for the torch_geometric version. I am using Ubuntu 22.04.3 LTS.

The packages in torch_geometric==2.3 is:

Package                  Version
------------------------ ----------
certifi                  2023.11.17
charset-normalizer       3.3.2
filelock                 3.13.1
fsspec                   2023.10.0
idna                     3.6
Jinja2                   3.1.2
joblib                   1.3.2
MarkupSafe               2.1.3
mpmath                   1.3.0
networkx                 3.2.1
numpy                    1.26.2
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.18.1
nvidia-nvjitlink-cu12    12.3.101
nvidia-nvtx-cu12         12.1.105
pip                      23.3.1
psutil                   5.9.6
pyparsing                3.1.1
requests                 2.31.0
scikit-learn             1.3.2
scipy                    1.11.4
setuptools               68.0.0
sympy                    1.12
threadpoolctl            3.2.0
torch                    2.1.1
torch-geometric          2.3.0
tqdm                     4.66.1
triton                   2.1.0
typing_extensions        4.8.0
urllib3                  2.1.0
wheel                    0.41.2

The environment for torch_geometric==2.4 is:

Package                  Version
------------------------ ----------
certifi                  2023.11.17
charset-normalizer       3.3.2
filelock                 3.13.1
fsspec                   2023.10.0
idna                     3.6
Jinja2                   3.1.2
joblib                   1.3.2
MarkupSafe               2.1.3
mpmath                   1.3.0
networkx                 3.2.1
numpy                    1.26.2
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.18.1
nvidia-nvjitlink-cu12    12.3.101
nvidia-nvtx-cu12         12.1.105
pip                      23.3.1
psutil                   5.9.6
pyparsing                3.1.1
requests                 2.31.0
scikit-learn             1.3.2
scipy                    1.11.4
setuptools               68.0.0
sympy                    1.12
threadpoolctl            3.2.0
torch                    2.1.1
torch_geometric          2.4.0
tqdm                     4.66.1
triton                   2.1.0
typing_extensions        4.8.0
urllib3                  2.1.0
wheel                    0.41.2

While running the script above using torch_geometric==2.3 it works fine, but when i am using torch_geometric==2.4 i get:

Traceback (most recent call last):
  File "/home/jovyan/ml-turing-gnn/dummy.py", line 65, in <module>
    net2.load_state_dict(state)
  File "/opt/pyenv/versions/3.11.4/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for Encoder:
        While copying the parameter named "layer1.convs.<gene___to___disease>.lin_l.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 5]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (5) at non-singleton dimension 1',).
        While copying the parameter named "layer1.convs.<gene___to___disease>.lin_r.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 5]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (5) at non-singleton dimension 1',).
        While copying the parameter named "layer1.convs.<disease___to___gene>.lin_l.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 5]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (5) at non-singleton dimension 1',).
        While copying the parameter named "layer1.convs.<disease___to___gene>.lin_r.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 5]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (5) at non-singleton dimension 1',).
        While copying the parameter named "layer2.convs.<gene___to___disease>.lin_l.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 32]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (32) at non-singleton dimension 1',).
        While copying the parameter named "layer2.convs.<gene___to___disease>.lin_r.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 32]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (32) at non-singleton dimension 1',).
        While copying the parameter named "layer2.convs.<disease___to___gene>.lin_l.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 32]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (32) at non-singleton dimension 1',).
        While copying the parameter named "layer2.convs.<disease___to___gene>.lin_r.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 32]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (32) at non-singleton dimension 1',).

Finally if i tried to load a file created using torch_geometric==2.3 on a torch_geometric==2.4 environment, I have the error i described befor.

Traceback (most recent call last):
  File "/home/jovyan/create_and_save_model.py", line 52, in <module>
    net2.load_state_dict(state)
  File "/home/jovyan/miniconda3/envs/pyg24/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for Encoder:
        Missing key(s) in state_dict: "layer1.convs.<v1___e0___v1>.lin_l.weight", "layer1.convs.<v1___e0___v1>.lin_l.bias", "layer1.convs.<v1___e0___v1>.lin_r.weight", "layer1.convs.<v0___e0___v1>.lin_l.weight", "layer1.convs.<v0___e0___v1>.lin_l.bias", "layer1.convs.<v0___e0___v1>.lin_r.weight", "layer1.convs.<v1___e0___v2>.lin_l.weight", "layer1.convs.<v1___e0___v2>.lin_l.bias", "layer1.convs.<v1___e0___v2>.lin_r.weight", "layer1.convs.<v2___e0___v0>.lin_l.weight", "layer1.convs.<v2___e0___v0>.lin_l.bias", "layer1.convs.<v2___e0___v0>.lin_r.weight", "layer1.convs.<v0___e0___v0>.lin_l.weight", "layer1.convs.<v0___e0___v0>.lin_l.bias", "layer1.convs.<v0___e0___v0>.lin_r.weight", "layer2.convs.<v1___e0___v1>.lin_l.weight", "layer2.convs.<v1___e0___v1>.lin_l.bias", "layer2.convs.<v1___e0___v1>.lin_r.weight", "layer2.convs.<v0___e0___v1>.lin_l.weight", "layer2.convs.<v0___e0___v1>.lin_l.bias", "layer2.convs.<v0___e0___v1>.lin_r.weight", "layer2.convs.<v1___e0___v2>.lin_l.weight", "layer2.convs.<v1___e0___v2>.lin_l.bias", "layer2.convs.<v1___e0___v2>.lin_r.weight", "layer2.convs.<v2___e0___v0>.lin_l.weight", "layer2.convs.<v2___e0___v0>.lin_l.bias", "layer2.convs.<v2___e0___v0>.lin_r.weight", "layer2.convs.<v0___e0___v0>.lin_l.weight", "layer2.convs.<v0___e0___v0>.lin_l.bias", "layer2.convs.<v0___e0___v0>.lin_r.weight". 
        Unexpected key(s) in state_dict: "layer1.convs.v0__e0__v0.lin_l.weight", "layer1.convs.v0__e0__v0.lin_l.bias", "layer1.convs.v0__e0__v0.lin_r.weight", "layer1.convs.v0__e0__v2.lin_l.weight", "layer1.convs.v0__e0__v2.lin_l.bias", "layer1.convs.v0__e0__v2.lin_r.weight", "layer1.convs.v1__e0__v0.lin_l.weight", "layer1.convs.v1__e0__v0.lin_l.bias", "layer1.convs.v1__e0__v0.lin_r.weight", "layer1.convs.v1__e0__v1.lin_l.weight", "layer1.convs.v1__e0__v1.lin_l.bias", "layer1.convs.v1__e0__v1.lin_r.weight", "layer1.convs.v2__e0__v1.lin_l.weight", "layer1.convs.v2__e0__v1.lin_l.bias", "layer1.convs.v2__e0__v1.lin_r.weight", "layer2.convs.v0__e0__v0.lin_l.weight", "layer2.convs.v0__e0__v0.lin_l.bias", "layer2.convs.v0__e0__v0.lin_r.weight", "layer2.convs.v0__e0__v2.lin_l.weight", "layer2.convs.v0__e0__v2.lin_l.bias", "layer2.convs.v0__e0__v2.lin_r.weight", "layer2.convs.v1__e0__v0.lin_l.weight", "layer2.convs.v1__e0__v0.lin_l.bias", "layer2.convs.v1__e0__v0.lin_r.weight", "layer2.convs.v1__e0__v1.lin_l.weight", "layer2.convs.v1__e0__v1.lin_l.bias", "layer2.convs.v1__e0__v1.lin_r.weight", "layer2.convs.v2__e0__v1.lin_l.weight", "layer2.convs.v2__e0__v1.lin_l.bias", "layer2.convs.v2__e0__v1.lin_r.weight". 

[EdisonLeeeee]

In your code example, the error is raised because FakeHeteroDataset generates heterogeneous graphs with different node feature dimensions and edge types. It is not related to the different PyG versions in this example. You might need to double-check your datasets to ensure that the dimensions or connections have not been changed.

[ver228]

Hey @EdisonLeeeee I apologised i made a couple of mistakes while copying the snippet. The observations I made still hold. Please take another look at it. I just edited it.

In short the mistakes I made are:

• I didn't copy the latest version of the code and skipped the most important part that is the load of the state dictionary with the an uninitialised model. This fails on pyg==2.4
• The error I copied while loading a 2.3 state saved model using 2.4 was caused by what you said on the FakeHeteroDataset but that was me being careless while copy the error since i saw it before in a fixed dataset, so i didn't put to much attention to what i was copying. The error is still there, the format between how the edge dictionaries are saved changed from pyg 2.3 to 2.4 and is breaking model saved on pyg 2.3 while read on pyg 2.4

[EdisonLeeeee]

Thank you for your inputs. I can now reproduce your error. This is because the recent update of ModuleDict in PyG has added additional < and > characters for the layer names, such as:

# in 2.3
layer1.convs.gene___to___disease 
# in 2.4
layer1.convs.<gene___to___disease>

It seems you need to manually add the < and > characters to the checkpoints saved by PyG 2.3. Sorry for the inconvenience :(

[ver228]

Sorry that's only one of the problems and it is not a major one. I agree it is not too difficult to edit the state dictionary.

The problem is really that in pyg==2.4 if the layers are initialised lazily as in the example above a the state dictionary of a pretrained model is not loaded correctly. It crashes with errors like:
While copying the parameter named "layer2.convs.<disease___to___gene>.lin_r.weight", whose dimensions in the model are torch.Size([0]) and whose dimensions in the checkpoint are torch.Size([32, 32]), an exception occurred : ('The size of tensor a (0) must match the size of tensor b (32) at non-singleton dimension 1',). Please try to run the example above, it will fail while loading net2 = Encoder(32, edges).

I guess one could do a forward pass into the uninitialised model, and then try to load the state dictionary, but it does defeat the purpose of having a lazy initialisation, plus this was not the behaviour in pyg=2.3. It seems to be a bug rather than a feature.

[rusty1s]

This should be already fixed on latest master. I need to push a hotfix release to integrate it.

[ver228]

Thanks! That's great :)

[AdrianSchneble]

This should be already fixed on latest master. I need to push a hotfix release to integrate it.

That's great to hear! I just had what seems to be the same problem as far as I can tell, and stumbled upon this discussion as I was about to create a new discussion.

Furthermore, I can verify that cloning the latest master version and installing it locally fixed the problem (at least for me).
Looking forward to the hotfix 👍