-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
72 lines (45 loc) · 2.14 KB
/
model.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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import torch
import torch.nn as nn
from arg import args
import torch.nn.functional as F
from torch_geometric.nn import GAE, VGAE, GATConv, BatchNorm, GCNConv
class prediction_MLP(nn.Module):
def __init__(self, in_dim, hidden_dim, out_dim): # bottleneck structure
super().__init__()
self.layer1 = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True)
)
self.layer2 = nn.Linear(hidden_dim, out_dim)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
return x
class GAE(torch.nn.Module):
def __init__(self, in_channels, hidden_num, head, out_channels, emb_size_1, emb_size_2, activation, base_model=GATConv):
super().__init__()
self.base_model = base_model
self.conv1 = base_model(in_channels, hidden_num, heads=head, dropout=args.dropout)
self.conv2 = base_model(hidden_num * head, out_channels, heads=1, concat=False,dropout=args.dropout)
self.bn1 = BatchNorm(hidden_num * head, momentum=0.99)
self.bn2 = BatchNorm(out_channels, momentum=0.99)
self.p_1 = prediction_MLP(out_channels, out_channels*2, emb_size_1)
self.p_2 = prediction_MLP(out_channels, out_channels*2, emb_size_2)
self.p_12 = prediction_MLP(out_channels, out_channels*2, (emb_size_1 + emb_size_2))
self.activation = activation
def encoder(self, x, edge_index):
x = F.dropout(x, p=args.dropout, training=self.training)
x = self.activation(self.bn1(self.conv1(x, edge_index)))
x = F.dropout(x, p=args.dropout, training=self.training)
x = self.activation(self.conv2(x, edge_index))
return x
def decoder_all(self, z):
recon_emb_1 = self.p_1(z)
recon_emb_2 = self.p_2(z)
recon_12 = self.p_12(z)
return recon_emb_1,recon_emb_2, recon_12
def forward(self, x: torch.Tensor, edge_index: torch.Tensor):
z = self.encoder(x, edge_index)
embs = self.decoder_all(z)
return embs