Using Edge weights in RGCN?

[Helaly96]

I am currently modelling a graph that has Different type of edges, let's say (X1) and (X2) and it can also have a w
For example we can say that Node A is connected to Node B with edge of type X1 AND that edge has weight of 10.

RGCNConv gives the ability to have the first part which is having a (W) matrix for each relation type, however i don't want to treat all edges equally.

how would you start implementing this?

[rusty1s]

If I understand you correctly, you are searching for a way to apply RGCNConv on weighted graphs? In that case, you can extend RGCNConv and add an edge_weight to its argument which is then used in message. Pseudo-code:

def forward(self, x, edge_index, edge_weight, edge_type):
    out = 0
    for i in range(num_edge_types):
        mask = edge_type == i
        out += self.propagate(edge_index[:, mask], edge_type=i, edge_weight=edge_weight[mask])
   return out
   
def message(self, x_j, edge_weight, edge_type):
    return edge_weight * x_j @ self.weight[edge_type]

[Helaly96]

class ExtendedRGCNConv(RGCNConv):
    def __init__(self, in_channels, out_channels, num_relations):
        super(ExtendedRGCNConv, self).__init__(in_channels, out_channels, num_relations)

    def forward(self, x, edge_index, edge_weight, edge_type):
        x_l= None
        if isinstance(x, tuple):
            x_l = x[0]
        else:
            x_l = x
        if x_l is None:
            x_l = torch.arange(self.in_channels_l, device=self.weight.device)

        x_r=x_l
        if isinstance(x, tuple):
            x_r = x[1]
        # propagate_type: (x: Tensor, edge_type_ptr: OptTensor)
        out = torch.zeros(x_r.size(0), self.out_channels, device=x_r.device)
        for i in range(self.num_relations):
            mask = edge_type == i
            if edge_weight is not None:
                out += self.propagate(edge_index[:, mask], x=x , edge_type=i, edge_weight=edge_weight[mask])
        return out
    def message(self, x_j, edge_weight, edge_type):
        return edge_weight * x_j @ self.weight[edge_type]

class RGCNModel(nn.Module):
    def __init__(self, num_node_features, num_classes, num_relations, hidden_channels):
        super(RGCNModel, self).__init__()

        self.conv1 = ExtendedRGCNConv(in_channels=num_node_features, out_channels=hidden_channels[1], num_relations=num_relations)
        self.conv2 = ExtendedRGCNConv(in_channels=hidden_channels[1], out_channels=num_classes, num_relations=num_relations)

    def forward(self, x, edge_index, edge_type,edge_attr):
        x = self.conv1(x=x, edge_index=edge_index, edge_type=edge_type,edge_weight=edge_attr)
        x = F.relu(x)
        x = self.conv2(x=x, edge_index=edge_index, edge_type=edge_type,edge_weight=edge_attr)
        return F.log_softmax(x, dim=1)

This works and by works i mean the model trains, however loss values exploded!
do you have any idea what caused this

[EdisonLeeeee]

Is your edge_weight normalized to a range of 0 to 1?

[Helaly96]

@EdisonLeeeee Thanks! After normalizing edge_weight, loss values are good again.

Also Thanks @rusty1s for your help.