wehub-resource-sync
41 lines
1.0 KiB
Python
41 lines
1.0 KiB
Python
import dgl
|
|
import torch as th
|
|
import torch.nn as nn
|
|
from DGLRoutingLayer import DGLRoutingLayer
|
|
from torch.nn import functional as F
|
|
|
|
g = dgl.DGLGraph()
|
|
g.graph_data = {}
|
|
|
|
in_nodes = 20
|
|
out_nodes = 10
|
|
g.graph_data["in_nodes"] = in_nodes
|
|
g.graph_data["out_nodes"] = out_nodes
|
|
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))
|
|
g.graph_data["in_indx"] = in_indx
|
|
g.graph_data["out_indx"] = out_indx
|
|
|
|
# add edges use edge broadcasting
|
|
for u in out_indx:
|
|
g.add_edges(in_indx, u)
|
|
# init states
|
|
f_size = 4
|
|
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)
|
|
|
|
routing_layer = DGLRoutingLayer(g)
|
|
|
|
entropy_list = []
|
|
for i in range(15):
|
|
routing_layer()
|
|
dist_matrix = g.edata["c"].view(in_nodes, out_nodes)
|
|
entropy = (-dist_matrix * th.log(dist_matrix)).sum(dim=0)
|
|
entropy_list.append(entropy.data.numpy())
|
|
std = dist_matrix.std(dim=0)
|