GraphKT/graph_net.py at main · abadrinath947/GraphKT · GitHub

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.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch import nn
from torch.nn import Linear, PReLU
from torch_geometric.nn import GCNConv, GATConv, SAGEConv

"""
class GCN(torch.nn.Module):
    def __init__(self, num_skills, hidden_dim):
        super().__init__()
        self.pre_embs = nn.Embedding(num_skills, hidden_dim)
        self.conv1 = GCNConv(hidden_dim, hidden_dim)
        self.prelu1 = PReLU()
        self.dropout = nn.Dropout(0.3)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)
        self.prelu2 = PReLU()
        self.conv3 = GCNConv(hidden_dim, hidden_dim)
        self.prelu3 = PReLU()
        self.out = Linear(hidden_dim, hidden_dim)

    def forward(self, x, edge_index, edge_weight):
        h1 = self.dropout(self.prelu1(self.conv1(self.pre_embs(x), edge_index, edge_weight = edge_weight)))
        h2 = self.dropout(self.prelu2(self.conv2(h1, edge_index, edge_weight = edge_weight)))
        h3 = self.prelu3(self.conv3(h2, edge_index, edge_weight = edge_weight))
        return self.out(h3)

        #return torch.cat([h1, h2, h3, h4], dim = -1)
"""

"""
class GCN(torch.nn.Module):
    def __init__(self, num_skills, hidden_dim):
        super().__init__()
        self.pre_embs = nn.Embedding(num_skills, hidden_dim)
        self.conv1 = GATConv(hidden_dim, hidden_dim)
        self.prelu1 = PReLU()
        self.dropout = nn.Dropout(0.3)
        self.conv2 = GATConv(hidden_dim, hidden_dim)
        self.prelu2 = PReLU()
        self.conv3 = GATConv(hidden_dim, hidden_dim)
        self.prelu3 = PReLU()
        self.out = Linear(hidden_dim, hidden_dim)

    def forward(self, x, edge_index, edge_weight):
        h1 = self.dropout(self.prelu1(self.conv1(self.pre_embs(x), edge_index)))
        h2 = self.dropout(self.prelu2(self.conv2(h1, edge_index)))
        #h3 = self.prelu3(self.conv3(h2, edge_index))
        return self.out(h2)
        #return torch.cat([h1, h2, h3, h4], dim = -1)
"""

class GCN(torch.nn.Module):
    def __init__(self, num_skills, hidden_dim):
        super().__init__()
        self.pre_embs = nn.Embedding(num_skills, hidden_dim)
        self.conv1 = SAGEConv(hidden_dim, hidden_dim)
        self.prelu1 = PReLU()
        self.dropout = nn.Dropout(0.3)
        self.conv2 = SAGEConv(hidden_dim, hidden_dim)
        self.prelu2 = PReLU()
        self.conv3 = SAGEConv(hidden_dim, hidden_dim)
        self.prelu3 = PReLU()
        self.out = Linear(hidden_dim, hidden_dim)

    def forward(self, x, edge_index, edge_weight):
        h1 = self.dropout(self.prelu1(self.conv1(self.pre_embs(x), edge_index)))
        h2 = self.dropout(self.prelu2(self.conv2(h1, edge_index)))
        #h3 = self.prelu3(self.conv3(h2, edge_index))
        return self.out(h2)
        #return torch.cat([h1, h2, h3, h4], dim = -1)