RGCN parameters: in_features, hidden_features, out_features


I aim to use a RGCN for link predicition so I am following: 5.3 Link Prediction — DGL 0.6.1 documentation

When getting to final step, I dont fully understand the following code:

def compute_loss(pos_score, neg_score):
# Margin loss
n_edges = pos_score.shape[0]
return (1 - pos_score.unsqueeze(1) + neg_score.view(n_edges, -1)).clamp(min=0).mean()

k = 5
model = Model(10, 20, 5, hetero_graph.etypes)
user_feats = hetero_graph.nodes[‘user’].data[‘feature’]
item_feats = hetero_graph.nodes[‘item’].data[‘feature’]
node_features = {‘user’: user_feats, ‘item’: item_feats}
opt = torch.optim.Adam(model.parameters())
for epoch in range(10):
negative_graph = construct_negative_graph(hetero_graph, k, (‘user’, ‘click’, ‘item’))
pos_score, neg_score = model(hetero_graph, negative_graph, node_features, (‘user’, ‘click’, ‘item’))
loss = compute_loss(pos_score, neg_score)

Those 10,20,5 in the Model arguments… where they come from? Checking the Model function parameters from the model defition I see they are in_features, hidden_features, out_features.

My question is: how can I guess those numbers for my graph?

As I said, my model has the same parameters as the tutorial (with a different graph, obviously).


in_features corresponds to the size of your input features.
out_features corresponds to the size of your output, usually the number of classes for classification or 1 for regression.
hidden_features corresponds to the size of your hidden state, where you set it as a hyperparameter.

1 Like

Thank you so much. One last thing. The number of classes, are in this cases the number of types of relations in the graph?

No, they are the number of output node classes for node classification. Different from the number of relations which tells how many edge types you have.

1 Like

But should i still use number of node classes if this is a link prediction task?

Oh I didn’t notice. In this case you don’t need the number of classes. You can set out_features to be the same as hidden_features so that the representations can be fed into the score predictor to produce the final positive/negative scores.