File size: 1,892 Bytes
bf620c6 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.ops import GCN, GraphUnet, Initializer, norm_g
class GNet(nn.Module):
def __init__(self, in_dim, n_classes, args):
super(GNet, self).__init__()
self.n_act = getattr(nn, args.act_n)()
self.c_act = getattr(nn, args.act_c)()
self.s_gcn = GCN(in_dim, args.l_dim, self.n_act, args.drop_n)
self.g_unet = GraphUnet(
args.ks, args.l_dim, args.l_dim, args.l_dim, self.n_act,
args.drop_n)
self.out_l_1 = nn.Linear(3*args.l_dim*(args.l_num+1), args.h_dim)
self.out_l_2 = nn.Linear(args.h_dim, n_classes)
self.out_drop = nn.Dropout(p=args.drop_c)
Initializer.weights_init(self)
def forward(self, gs, hs, labels):
hs = self.embed(gs, hs)
logits = self.classify(hs)
return self.metric(logits, labels)
def embed(self, gs, hs):
o_hs = []
for g, h in zip(gs, hs):
h = self.embed_one(g, h)
o_hs.append(h)
hs = torch.stack(o_hs, 0)
return hs
def embed_one(self, g, h):
g = norm_g(g)
h = self.s_gcn(g, h)
hs = self.g_unet(g, h)
h = self.readout(hs)
return h
def readout(self, hs):
h_max = [torch.max(h, 0)[0] for h in hs]
h_sum = [torch.sum(h, 0) for h in hs]
h_mean = [torch.mean(h, 0) for h in hs]
h = torch.cat(h_max + h_sum + h_mean)
return h
def classify(self, h):
h = self.out_drop(h)
h = self.out_l_1(h)
h = self.c_act(h)
h = self.out_drop(h)
h = self.out_l_2(h)
return F.log_softmax(h, dim=1)
def metric(self, logits, labels):
loss = F.nll_loss(logits, labels)
_, preds = torch.max(logits, 1)
acc = torch.mean((preds == labels).float())
return loss, acc
|