Batches creation for graphormer layer (transformers layer)
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Some utilities to convert the dgl.batch graph into something the graphormer layer can work with (or revisiting the graphormer layer).
## Motivation
In my opinion, one of the strengths of DGL is the simple way in which you can create batches, dataloaders and work with different architectures easily, but, with the graphormer layer, you have to change the batch into a different format.
## Alternatives
I think the possible solutions are:
- add some utilities to convert the classical dgl.batch (perfect for a message-passing gnn) to a transformer batch
- change the graphormer layer to work with the dgl.batch
## Pitch
I'm working on a graph classification task and would like to go from GAT (for instance) to Graphormer as smooth as possible.
## Additional context
This is the forward of my network (currently I'm trying to improve this code):
```python
def forward(self, g, nfeats):
# graphormer layer needs features (batch_size, N, N, dim_features), N is the maximum number of nodes that i can find in a graph in the current batch, and, attention_mask (batch_size, N, N)
g_list = dgl.unbatch(g)
max_num_nodes = torch.max(g.batch_num_nodes())
features = torch.zeros(len(g_list), max_num_nodes.item(), self.num_features).to(device)
attn_mask = torch.zeros(len(g_list), max_num_nodes.item(), max_num_nodes.item()).to(device)
batch_n = 0
accum_nodes = 0
for graph in g_list:
features[batch_n, :graph.number_of_nodes()] = nfeats[accum_nodes:accum_nodes+graph.number_of_nodes()] + self.centrality_encoder(graph)
attn_mask[batch_n, :, graph.number_of_nodes():] = 1
attn_mask[batch_n, graph.number_of_nodes():, 1:] = 1
batch_n += 1
accum_nodes += graph.number_of_nodes()
# compute spatial encoding (it works with a batched dgl graphs)
bias = self.encoder(g)
# apply graphormer layer
h = self.transformer(features, bias, attn_mask=attn_mask)
hg = torch.zeros(len(g_list), self.num_features).to(device)
batch_n = 0
# here, instead of the mean, i can use the virtual node
for graph in g_list:
hg[batch_n] = torch.mean(h[batch_n, :graph.number_of_nodes(), :], dim=0)
batch_n += 1
hg = F.dropout(hg, p=0.2, training=self.training)
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.