-
Notifications
You must be signed in to change notification settings - Fork 0
/
embedder.py
36 lines (27 loc) · 1.06 KB
/
embedder.py
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
import numpy as np
import torch
import torch.nn as nn
from torch_geometric.nn import GCNConv
# To fix the random seed
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False
class embedder:
def __init__(self, args):
self.args = args
self.hidden_layers = eval(args.layers)
class Encoder(nn.Module):
def __init__(self, layer_config):
super().__init__()
self.stacked_gnn = nn.ModuleList(
[GCNConv(layer_config[i - 1], layer_config[i]) for i in range(1, len(layer_config))])
self.stacked_bns = nn.ModuleList(
[nn.BatchNorm1d(layer_config[i], momentum=0.01) for i in range(1, len(layer_config))])
self.stacked_prelus = nn.ModuleList([nn.PReLU() for _ in range(1, len(layer_config))])
def forward(self, x, edge_index):
for i, gnn in enumerate(self.stacked_gnn):
x = gnn(x, edge_index, edge_weight=None)
x = self.stacked_bns[i](x)
x = self.stacked_prelus[i](x)
return x