MLP with heterograph

[ogggcar]

Hi everyone.

I have an heterograph with this metadata: metagraph=[('person', 'thing', 'link1'), ('person', 'person', 'link2')])

I’m trying to use an MLP to generate edge specific embeddings for a given node. This is the predictor part of my model:

class MLP(nn.Module):
    def __init__(self, out_feats):
        super().__init__()
        self.layer = nn.Sequential(
                        nn.Linear(out_feats, out_feats),
                        nn.ReLU(), 
                        nn.Linear(out_feats, out_feats)
                     )
    def forward(self, x):
        return self.layer(x)

class HeteroDotProductPredictor(nn.Module):
    def __init__(self, out_feats):
        super().__init__()
        self.etype_project = {('person', 'link1', 'thing'): MLP(out_feats), 
                              ('person', 'link2', 'person'): MLP(out_feats)}

    def forward(self, graph, h, etype):
        with graph.local_scope():
            graph.ndata['h'] = self.etype_project[etype](h)
            graph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype)
            return graph.edges[etype].data['score']

However, I get the following error in the training loop:

TypeError: linear(): argument 'input' (position 1) must be Tensor, not dict

I understand this is because my MLP need a Tensor and not the dict, but since I have two type of nodes, I’m not sure about how could I pass the dict through the MLP.

Thanks everyone.

[BarclayII]

Looks like you want to compute the person features with an MLP and thing features with another MLP, and then compute a dot product? If so, since h is a dict of node types and features, you need to run your MLP on each node type separately, something like:

graph.nodes['person'].data['h'] = self.ntype_project['person'](h['person'])