Discrepancy in HeteroData.node_offsets When Using PyG Model Directly vs. Recording with TensorBoard's SummaryWriter

[Julian-Yu]

When debugging a model created with PyG, if I use y=model(current_time) directly the breakpoints show HeteroData.node_offsets as {'0': 0, '1': 4, '2': 6}, which is what I expect to get as output. But when I record the neural network structure using SummaryWriter in torch.utils.tensorboard, using w.add_graph(model, current_time), at the same breakpoint it shows HeteroData.node_offsets as {'0': 0, '1' : tensor(10), '2': tensor(10)}. May I ask why this is?

[rusty1s]

Do you have a small example for me to reproduce?

[Julian-Yu]

Thank you for writing back. Here's my small example:

import torch
import torch_geometric as pyg
from torch.utils.tensorboard import SummaryWriter

h = {'0': {'x': torch.tensor([[1., 0.],
                              [2., 1.],
                              [3., 2.],
                              [4., 3.]]),
           'id': torch.tensor([0, 1, 2, 3], dtype=torch.int32)},
     '1': {'x': torch.tensor([[1., 1.],
                              [2., 3.]]),
           'id': torch.tensor([4, 5], dtype=torch.int32)},
     '2': {'x': torch.tensor([[2., 0.],
                              [3., 1.],
                              [4., 2.],
                              [5., 3.]]),
           'id': torch.tensor([6, 7, 8, 9], dtype=torch.int32)},
     ('0', '0', '1'): {'edge_index': torch.tensor([[0, 1, 2, 3],
                                                   [4, 5, 4, 5]], dtype=torch.int32)},
     ('1', '0', '0'): {'edge_index': torch.tensor([[4, 5, 4, 5],
                                                   [0, 1, 2, 3]], dtype=torch.int32)},
     ('1', '1', '2'): {'edge_index': torch.tensor([[4, 5, 4, 5],
                                                   [6, 7, 8, 9]], dtype=torch.int32)},
     ('2', '1', '1'): {'edge_index': torch.tensor([[6, 7, 8, 9],
                                                   [4, 5, 4, 5]], dtype=torch.int32)}}

hg = pyg.data.HeteroData.from_dict(h)

def absolute_index2relative_index(absolute_index: dict[tuple[str, str, str], torch.Tensor],
                                  node_offsets: dict[str, int]) -> dict[tuple[str, str, str], torch.Tensor]:
    """
    Converting Absolute Indexes to Relative Indexes in HeteroData
    :param node_offsets: Offsets for different types of nodes
    :param absolute_index: absolute_index
    :return: relative_index
    """
    relative_index = absolute_index
    for edge_type, value in absolute_index.items():
        src_type, dst_type = edge_type[0], edge_type[2]
        value[0, :] = value[0, :] - node_offsets[src_type]
        value[1, :] = value[1, :] - node_offsets[dst_type]
    return relative_index

class MyConv(torch.nn.Module):
    def __init__(self,
                 hg: pyg.data.HeteroData,
                 in_channels: int,
                 out_channels: int):
        super(MyConv, self).__init__()
        self.hg = hg
        self.hanlayer = pyg.nn.HANConv(in_channels, out_channels, num_heads=2, metadata=hg.metadata())
        self.lin = torch.nn.Linear(out_channels, 1)

    def forward(self, a: torch.Tensor) -> torch.Tensor:
        # Since w.add_graph must be passed an input of type torch.Tensor,
        # the "a" is only used as a placeholder for the input and is not actually passed into the model.
        edge_index = absolute_index2relative_index(self.hg.edge_index_dict, self.hg.node_offsets)
        # -> We can add a breakpoint below this line to see the value of self.hg.node_offsets
        x = self.hanlayer(self.hg.x_dict, edge_index)
        return self.lin(x)

model = MyConv(hg, 2, 2)

with SummaryWriter(log_dir="E:/") as w:
    w.add_graph(model, torch.tensor([[1., 0.],
                                     [2., 1.],
                                     [3., 2.],
                                     [4., 3.]]))

When we insert a breakpoint at the place of the comment, we can see that the value of self.hg.node_offsets is
{'0': 0, '1': tensor(10), '2': tensor(10)},
and from the data of "hg", we can easily deduce that the correct output should be
{ '0': 0, '1': 4, '2': 6}
As shown in the figure below:

Thank you very much for your help!

[rusty1s]

Thanks. Fixed in #8430.