Create your customized graph construction method
In this subsection we give the example for model VGG-16 as computational graph use Torch-Geometric package. Create the graph by convolutional layer's in and out channels. Write a class for converting DNN to computational graph:
from gnnrl.graph_env.graph_construction import conv_motif,conv_sub_graph
from gnnrl.utils.batchwise_graphs import get_next_graph_batch
import torch
from torch_geometric.data import Data
from torch_geometric.data import DataLoader
class computational_graph_pyg():
def __init__(self, in_channels, out_channels, feature_size, device=None, pruning_ratios=1):
self.in_channels = in_channels
self.out_channels = out_channels
self.pruning_ratios = pruning_ratios
self.c_in_channels = self.in_channels[1:] * pruning_ratios
self.c_out_channels = self.out_channels[:-1] * pruning_ratios
self.feature_size = feature_size
if device is None:
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
self.device = device
self.plain_graph = self.plain_computational_graph()
self.hierachical_graph = self.hierarchical_computational_graph()
def plain_computational_graph(self):
return None
def hierarchical_computational_graph(self):
hierarchical_graph={}
hierarchical_graph['level1'] = self.level1_graph().to(self.device)
hierarchical_graph['level2'] = self.level2_graph().to(self.device)
return hierarchical_graph
def level1_graph(self):
level_1_graphs = []
for in_c in self.c_in_channels:
edge_index = conv_motif(in_c)
G = Data(edge_index=torch.tensor(edge_index).long().t().contiguous())
G.x = torch.randn([G.num_nodes,self.feature_size]).to(self.device)
level_1_graphs.append(G.to(self.device))
level_1_graphs = DataLoader(level_1_graphs,batch_size=len(level_1_graphs), shuffle=False)
level_1_graphs = get_next_graph_batch(level_1_graphs)
return level_1_graphs
def level2_graph(self):
node_cur = 0
edge_list = []
edge_type = []
k = 0 # layer index
normal_ope_edge_type = len(self.c_out_channels)
for i in range(len(self.c_out_channels)):
edge_list,edge_type,node_cur = conv_sub_graph(node_cur,self.c_out_channels[i],edge_list,edge_type,i,normal_ope_edge_type)
k+=1
#Batch Norm
edge_list.append([node_cur,node_cur+1])
edge_type.append(normal_ope_edge_type)
node_cur += 1
Graph = Data(edge_index=torch.tensor(edge_list).t().contiguous(),edge_type =edge_type)
Graph.x = torch.randn([Graph.num_nodes, self.feature_size])
Graph.edge_features = None
Graph = DataLoader([Graph],batch_size=1, shuffle=False)
Graph = get_next_graph_batch(Graph)
return Graph
def update_pruning_ratio(self,pruning_ratios):
self.pruning_ratios = pruning_ratios
self.c_in_channels = self.in_channels[1:] * pruning_ratios
self.c_out_channels = self.out_channels[:-1] * pruning_ratios
self.plain_graph = self.plain_computational_graph()
self.hierachical_graph = self.hierarchical_computational_graph()
Then load the neural network from torchvision.models
,
from torchvision.models import vgg16
net = vgg16()
Next, get the network's convolutional layers' input and output channels,
import torch.nn as nn
in_channels = []
out_channels=[]
for name,layer in net.named_modules():
if isinstance(layer,nn.Conv2d):
in_channels.append(layer.in_channels)
out_channels.append(layer.out_channels)
Initialize the computational_graph_pyg
object you created,
pyg_g = computational_graph_pyg(in_channels,out_channels,feature_size=20)
Get your computational graph by calling the function plain_computational_graph()
,
graph = pyg_g.plain_computational_graph()