Lazy initialisation and torch functions

[MattBortoletto]

Hi everyone,

I'm building a heterogeneous GNN by converting a homogeneous one using to_hetero or to_hetero_with_bases. Then I would like to see the difference in the number of parameters between the two versions. Therefore I define a simple function for counting parameters:

def count_parameters(model, data):
    model.train()
    model(data.x_dict, data.edge_index_dict)
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

The first two lines inside the functions are to initialise the model parameter, as the GNN class uses lazy initialisation. When I pass the GNN obtained using to_hetero it runs fine, but when I pass the one obtained with to_hetero_with_bases I get the same error as if I didn't initialise the model:

ValueError: Attempted to use an uninitialized parameter in <method 'numel' of 'torch._C._TensorBase' objects>. This error happens when you are using a `LazyModule` or explicitly manipulating `torch.nn.parameter.UninitializedParameter` objects. When using LazyModules Call `forward` with a dummy batch to initialize the parameters before calling torch functions

Is it a bug or am I missing something? Is there a smarter way to check the number of parameters?

Thank you for your answers!

[rusty1s]

Looks like a bug TBH. Do you have a simple example to reproduce this? How does your model look like?

[MattBortoletto]

Here's a small example:

import torch
import torch.nn.functional as F
from torch_geometric.data import HeteroData
from torch_geometric.nn import GATConv, Linear, to_hetero, to_hetero_with_bases, global_mean_pool

x = torch.randn(5, 12)
top = torch.randint(0, 5, (2, 5))
bottom = torch.randint(0, 5, (2, 5))
left = torch.randint(0, 5, (2, 5))
right = torch.randint(0, 5, (2, 5))

data = HeteroData()
data['obj', 'top', 'obj'].edge_index = top
data['obj', 'bottom', 'obj'].edge_index = bottom
data['obj', 'left', 'obj'].edge_index = left
data['obj', 'right', 'obj'].edge_index = right
data['obj'].x = x

def count_parameters(model, data):
    model.train()
    model(data.x_dict, data.edge_index_dict)
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        # Readout layer
        x = global_mean_pool(x, batch=None)
        return x

model = GAT(hidden_channels=64, out_channels=64)
model_basis = to_hetero_with_bases(model, data.metadata(), num_bases=3)
count_parameters(model, data)

[rusty1s]

Sorry for the late reply - I was on vacation. I sadly cannot reproduce this issue - which PyG version are you using? It might have been already fixed in master. Added some tests to confirm: #5185