1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
|
import torch
from ogb.utils.features import get_bond_feature_dims
from torch_geometric.graphgym.register import (
register_edge_encoder,
register_node_encoder,
)
@register_node_encoder('example')
class ExampleNodeEncoder(torch.nn.Module):
"""Provides an encoder for integer node features.
Args:
num_classes (int): The number of classes for the embedding mapping to
learn.
"""
def __init__(self, emb_dim, num_classes=None):
super().__init__()
self.encoder = torch.nn.Embedding(num_classes, emb_dim)
torch.nn.init.xavier_uniform_(self.encoder.weight.data)
def forward(self, batch):
# Encode just the first dimension if more exist
batch.x = self.encoder(batch.x[:, 0])
return batch
@register_edge_encoder('example')
class ExampleEdgeEncoder(torch.nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.bond_embedding_list = torch.nn.ModuleList()
full_bond_feature_dims = get_bond_feature_dims()
for i, dim in enumerate(full_bond_feature_dims):
emb = torch.nn.Embedding(dim, emb_dim)
torch.nn.init.xavier_uniform_(emb.weight.data)
self.bond_embedding_list.append(emb)
def forward(self, batch):
bond_embedding = 0
for i in range(batch.edge_feature.shape[1]):
bond_embedding += \
self.bond_embedding_list[i](batch.edge_attr[:, i])
batch.edge_attr = bond_embedding
return batch
|
