-
Notifications
You must be signed in to change notification settings - Fork 4
/
train.py
67 lines (50 loc) · 1.59 KB
/
train.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
# -*- coding: utf-8 -*-
import torch
import torch.nn.functional as F
import networkx as nx
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from model import GCN
from graph_builder import build_karate_club_graph
import warnings
warnings.filterwarnings('ignore')
net = GCN(34, 5, 2)
print(net)
G = build_karate_club_graph()
inputs = torch.eye(34)
labeled_nodes = torch.tensor([0, 33]) # only the instructor and the president nodes are labeled
labels = torch.tensor([0, 1])
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
all_logits = []
for epoch in range(20):
logits = net(G, inputs)
all_logits.append(logits.detach())
logp = F.log_softmax(logits, 1)
# compute loss for labeled nodes
loss = F.nll_loss(logp[labeled_nodes], labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch %d | Loss: %.4f' % (epoch, loss.item()))
def draw(i):
cls1color = '#00FFFF'
cls2color = '#FF00FF'
pos = {}
colors = []
for v in range(34):
pos[v] = all_logits[i][v].numpy()
cls = pos[v].argmax()
colors.append(cls1color if cls else cls2color)
# ax.cla()
# ax.axis('off')
# ax.set_title('Epoch: %d' % i)
nx.draw_networkx(nx_G.to_undirected(), pos, node_color=colors, with_labels=True, node_size=300)
nx_G = G.to_networkx().to_undirected()
print(G.to_networkx())
# fig, ax = plt.subplots()
ax = plt.figure(dpi=150)
# fig.clf()
# ax = fig.subplots()
# draw(19) # draw the prediction of the first epoch
ani = animation.FuncAnimation(ax, draw, frames=len(all_logits), interval=200)
plt.show()