When using edge_weight with SAGEConv, I encountered ValueError due to unexpected tensor size

[yuan-meng-dd]

I have a link prediction model that works when edge_weight = None but throws ValueError: Encountered tensor with size 1814 in dimension -2, but expected size 10.0. when I supply actual edge weights to SAGEConv. I was able to reproduce this error on MovieLens data (code adapted from tutorial). In this case, edge_weight is the star rating given by a user to a movie.

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()

        self.conv1 = SAGEConv(hidden_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)

    def forward(self, x: Tensor, edge_index: Tensor, edge_weight: Tensor) -> Tensor:
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = self.conv2(x, edge_index, edge_weight)
        return x

class Classifier(torch.nn.Module):
    def forward(
        self, x_user: Tensor, x_movie: Tensor, edge_label_index: Tensor
    ) -> Tensor:
        # convert node embeddings to edge-level representations:
        edge_feat_user = x_user[edge_label_index[0]]
        edge_feat_movie = x_movie[edge_label_index[1]]

        # apply dot-product to get a prediction per supervision edge:
        return (edge_feat_user * edge_feat_movie).sum(dim=-1)

class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()

        self.movie_lin = torch.nn.Linear(data["movie"].num_features, hidden_channels)

        self.user_emb = torch.nn.Embedding(data["user"].num_nodes, hidden_channels)
        self.movie_emb = torch.nn.Embedding(data["movie"].num_nodes, hidden_channels)

        # instantiate homogeneous GNN:
        self.gnn = GNN(hidden_channels)

        # convert GNN model into a heterogeneous variant:
        self.gnn = to_hetero(self.gnn, metadata=data.metadata())

        # final classification layer
        self.classifier = Classifier()

    def forward(self, data: HeteroData) -> Tensor:
        x_dict = {
            "user": self.user_lin(data["user"].x),
            "movie": self.movie_lin(data["movie"].x)
            + self.movie_emb(data["movie"].node_id),
        }

        x_dict = self.gnn(x_dict, data.edge_index_dict, data.edge_weight_dict)

        pred = self.classifier(
            x_dict["user"],
            x_dict["movie"],
            data["user", "rates", "movie"].edge_label_index,
        )

        # return logit output for link prediction
        return pred

When I run the model, I encounter the aforementioned error:

# instantiate model
model = Model(hidden_channels=64)

# train for some epochs
N_EPOCHS = 5
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0002)

for epoch in range(N_EPOCHS):
    total_loss = total_examples = 0
    for sampled_data in tqdm.tqdm(train_loader):
        # reset gradients to 0
        optimizer.zero_grad()

        # get a batch of data
        sampled_data.to(device)
        
        # link prediction (logits)
        pred = model(sampled_data)
        
        # link labels (0 or 1)
        ground_truth = sampled_data["user", "rates", "movie"].edge_label
        
        # BCE loss 
        loss = F.binary_cross_entropy_with_logits(pred, ground_truth)
        
        loss.backward(retain_graph=True)
        optimizer.step()
        total_loss += float(loss) * pred.numel()
        total_examples += pred.numel()

    # save model
    torch.save(model, f"/dbfs{model_output_path}epoch={str(epoch)}.pt")
    print(f"epoch: {epoch:03d}, loss: {total_loss / total_examples:.4f}")

x_dict = self.gnn(x_dict, data.edge_index_dict, data.edge_weight_dict) seems to be the offensive line:

File <eval_with_key>.39:18, in forward(self, x, edge_index, edge_weight)
     16 relu__user = torch.nn.functional.relu(conv1__user, inplace = False);  conv1__user = None
     17 relu__movie = torch.nn.functional.relu(conv1__movie, inplace = False);  conv1__movie = None
---> 18 conv2__movie = self.conv2.user__intended_for__movie((relu__user, relu__movie), edge_index__user__intended_for__movie, edge_weight__user__intended_for__movie);  edge_index__user__intended_for__movie = edge_weight__user__intended_for__movie = None
     19 conv2__user = self.conv2.movie__intended_by__user((relu__movie, relu__user), edge_index__movie__intended_by__user, edge_weight__movie__intended_by__user);  relu__movie = relu__user = edge_index__movie__intended_by__user = edge_weight__movie__intended_by__user = None
     20 return {'user': conv2__user, 'movie': conv2__movie}

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/torch/nn/modules/module.py:1190, in Module._call_impl(self, *input, **kwargs)
   1186 # If we don't have any hooks, we want to skip the rest of the logic in
   1187 # this function, and just call forward.
   1188 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1189         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1190     return forward_call(*input, **kwargs)
   1191 # Do not call functions when jit is used
   1192 full_backward_hooks, non_full_backward_hooks = [], []

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-40c5be2e-59ec-4a21-8d8c-eed513799743/lib/python3.9/site-packages/torch_geometric/nn/conv/sage_conv.py:131, in SAGEConv.forward(self, x, edge_index, size)
    128     x = (self.lin(x[0]).relu(), x[1])
    130 # propagate_type: (x: OptPairTensor)
--> 131 out = self.propagate(edge_index, x=x, size=size)
    132 out = self.lin_l(out)
    134 x_r = x[1]

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-40c5be2e-59ec-4a21-8d8c-eed513799743/lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py:456, in MessagePassing.propagate(self, edge_index, size, **kwargs)
    453     for arg in decomp_args:
    454         kwargs[arg] = decomp_kwargs[arg][i]
--> 456 coll_dict = self.__collect__(self.__user_args__, edge_index,
    457                              size, kwargs)
    459 msg_kwargs = self.inspector.distribute('message', coll_dict)
    460 for hook in self._message_forward_pre_hooks.values():

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-40c5be2e-59ec-4a21-8d8c-eed513799743/lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py:328, in MessagePassing.__collect__(self, args, edge_index, size, kwargs)
    326     assert len(data) == 2
    327     if isinstance(data[1 - dim], Tensor):
--> 328         self.__set_size__(size, 1 - dim, data[1 - dim])
    329     data = data[dim]
    331 if isinstance(data, Tensor):

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-40c5be2e-59ec-4a21-8d8c-eed513799743/lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py:242, in MessagePassing.__set_size__(self, size, dim, src)
    240     size[dim] = src.size(self.node_dim)
    241 elif the_size != src.size(self.node_dim):
--> 242     raise ValueError(
    243         (f'Encountered tensor with size {src.size(self.node_dim)} in '
    244          f'dimension {self.node_dim}, but expected size {the_size}.'))

ValueError: Encountered tensor with size 1814 in dimension -2, but expected size 10.0.

If this is not the correct usage of SAGEConv, may I ask how I correctly can use egde_weight? Thanks!

[EdisonLeeeee]

edge_weight is not a supported argument in SAGEConv, you can use GraphConv instead.

[yuan-meng-dd]

I see. I was confused because I saw the edge_weight argument in DGL's doc and thought it was supported now. I didn't realize the PyG implementation is different.