Complex-Network-Project/model_processer.py at master · ShipXu/Complex-Network-Project · GitHub

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
from utils import Node
from utils import Digraph
from keras.utils import plot_model
from models.lenet import Lenet
from models.resnet50 import ResNet50
from models.mobilenet import MobileNet
from models.nasnet import NASNet

def visualize_model(model, path='models_visualization/model.png'):
    plot_model(model, to_file=path)

def read_models(model):
    top_tensors = {}
    bottom_tensors = {}
    layer_map = {}
    for layer in model.layers:
        layer_map[layer.name] = layer
        for i in range(len(layer._inbound_nodes)):
            outTensor = layer.get_output_at(i)
            if not isinstance(outTensor, list):
                if outTensor not in top_tensors:
                    top_tensors[outTensor] = [layer]
                else:
                    top_tensors[outTensor].append(layer)
            else:
                for tensor in outTensor:
                    if tensor not in top_tensors:
                        top_tensors[tensor] = [layer]
                    else:
                        top_tensors[tensor].append(layer)

        for i in range(len(layer._inbound_nodes)):
            inTensor = layer.get_input_at(i)
            if not isinstance(inTensor, list):
                if inTensor not in bottom_tensors:
                    bottom_tensors[inTensor] = [layer]
                else:
                    bottom_tensors[inTensor].append(layer)
            else:
                for tensor in inTensor:
                    if tensor not in bottom_tensors:
                        bottom_tensors[tensor] = [layer]
                    else:
                        bottom_tensors[tensor].append(layer)
    return layer_map, top_tensors, bottom_tensors

def model_to_dag(model):
    dag = Digraph()
    _, top_tensors, bottom_tensors = read_models(model)
    for key in top_tensors:
        if key in bottom_tensors:
            for from_layer in top_tensors[key]:
                for top_layer in bottom_tensors[key]:
                    if from_layer != top_layer:
                        from_node = Node(from_layer.name, from_layer)
                        top_node = Node(top_layer.name, top_layer)
                        dag.addEdge(from_node, top_node)
    return dag

if __name__ == '__main__':
    # model = Lenet()
    # model = ResNet50()
    # model = MobileNet()
    model = NASNet()
    visualize_model(model, path='models_visualization/NASNet.png')
    dag = model_to_dag(model)
    print(dag)