-
Notifications
You must be signed in to change notification settings - Fork 0
/
02.alexnet.frontend.py
48 lines (38 loc) · 1.75 KB
/
02.alexnet.frontend.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
import os,sys
from partition.interface import *
if __name__ == '__main__':
output_dirs= './models'
if not os.path.exists(output_dirs):
# Create a new directory because it does not exist
os.makedirs(output_dirs)
print("The output directory %s is created!" % (output_dirs))
origin_model = "bvlcalexnet-9.onnx"
### This template model Can be downloaded from:
# https://github.com/onnx/models/blob/main/vision/classification/alexnet/model/bvlcalexnet-9.onnx
input_model = format_onnx(origin_model)
model = onnx.load(input_model)
# resourceid = { 1:'lenovo_cpu0', 2:'lenovo_cpu1'}
### How the layers in alexnet model distribute over two computing nodes???
# random_map = {
# "lenovo_cpu0":
# ["conv1_1", "conv1_2", "norm1_1", "pool1_1", "conv2_1",
# "conv2_2", "norm2_1", "pool2_1", "conv3_1", "conv3_2",
# "conv4_1", "conv4_2", "conv5_1", "conv5_2", "pool5_1",
# "OC2_DUMMY_0", "fc6_1", "fc6_2"],
# "lenovo_cpu1":
# ["fc6_3", "fc7_1", "fc7_2", "fc7_3", "fc8_1", "prob_1"]
# }
platforms = ['lenovo'] # platforms = ['nx01', 'nx02']
random_map = load_json("mapping.json")
InputSpecs = Interface(model=input_model, mappings=random_map, platforms=platforms)
#cppname, NodesList, ComputingNodes
### Generate Cpp Files for Multinode-Inference.
GenerateCode = EngineCode(
CppName = "./models/multinode",
Platforms = InputSpecs.platforms,
NodesList = InputSpecs.nodes,
ComputingNodes = InputSpecs.computingnodes,
ValueInfos = InputSpecs.value_map,
Inputs = InputSpecs.inputs,
Outputs = InputSpecs.outputs,
Benchmark = False)