simple_routing.py 1.06 KB
Newer Older
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1
import dgl
2
3
4
import torch as th
import torch.nn as nn
from DGLRoutingLayer import DGLRoutingLayer
5
from torch.nn import functional as F
6
7
8
9
10
11

g = dgl.DGLGraph()
g.graph_data = {}

in_nodes = 20
out_nodes = 10
12
13
g.graph_data["in_nodes"] = in_nodes
g.graph_data["out_nodes"] = out_nodes
14
15
16
17
18
19
all_nodes = in_nodes + out_nodes
g.add_nodes(all_nodes)


in_indx = list(range(in_nodes))
out_indx = list(range(in_nodes, in_nodes + out_nodes))
20
21
g.graph_data["in_indx"] = in_indx
g.graph_data["out_indx"] = out_indx
22
23
24
25
26
27
28

# add edges use edge broadcasting
for u in out_indx:
    g.add_edges(in_indx, u)

# init states
f_size = 4
29
30
31
g.ndata["v"] = th.zeros(all_nodes, f_size)
g.edata["u_hat"] = th.randn(in_nodes * out_nodes, f_size)
g.edata["b"] = th.randn(in_nodes * out_nodes, 1)
32
33
34

routing_layer = DGLRoutingLayer(g)

35
entropy_list = []
36
37
for i in range(15):
    routing_layer()
38
    dist_matrix = g.edata["c"].view(in_nodes, out_nodes)
39
40
41
    entropy = (-dist_matrix * th.log(dist_matrix)).sum(dim=0)
    entropy_list.append(entropy.data.numpy())
    std = dist_matrix.std(dim=0)