onnx图构建

import onnx
from onnx import helper
from onnx import TensorProto
import numpy as np
from onnx import shape_inference


# layer-1   H_dim
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 224, 224])
weight_1 = helper.make_tensor('weight_1', TensorProto.FLOAT, [64, 3, 7, 7], np.random.randn(64, 3, 7, 7))
bias_1 = helper.make_tensor('bias_1', TensorProto.FLOAT, [64], np.random.randn(64))
conv_1 = helper.make_tensor_value_info('conv_1', TensorProto.FLOAT, [1, 64, 112, 112])
relu_1 = helper.make_tensor_value_info('relu_1', TensorProto.FLOAT, [1, 64, 112, 112])
node_conv_1 = helper.make_node(
    'Conv',
    ['input', 'weight_1', 'bias_1'],
    ['conv_1'],
    kernel_shape=[7,7],
    pads=[3,3,3,3],
    strides=[2,2],
)
node_relu_1 = helper.make_node(
    'Relu',
    ['conv_1'],
    ['relu_1'],
)

# layer-2   C_dim
weight_2 = helper.make_tensor('weight_2', TensorProto.FLOAT, [3, 64, 1, 1], np.random.randn(3, 64, 1, 1))
bias_2 = helper.make_tensor('bias_2', TensorProto.FLOAT, [3], np.random.randn(3))
conv_2 = helper.make_tensor_value_info('conv_2', TensorProto.FLOAT, [1, 3, 112, 112])
relu_2 = helper.make_tensor_value_info('relu_2', TensorProto.FLOAT, [1, 3, 112, 112])
node_conv_2 = helper.make_node(
    'Conv',
    ['relu_1', 'weight_2', 'bias_2'],
    ['conv_2'],
    kernel_shape=[1,1],
    pads=[0,0,0,0],
    strides=[1,1],
)
node_relu_2 = helper.make_node(
    'Relu',
    ['conv_2'],
    ['relu_2'],
)

# layer-3   H_dim
weight_3 = helper.make_tensor('weight_3', TensorProto.FLOAT, [16, 3, 3, 3], np.random.randn(16, 3, 3, 3))
bias_3 = helper.make_tensor('bias_3', TensorProto.FLOAT, [16], np.random.randn(16))
conv_3 = helper.make_tensor_value_info('conv_3', TensorProto.FLOAT, [1, 16, 112, 112])
relu_3 = helper.make_tensor_value_info('relu_3', TensorProto.FLOAT, [1, 16, 112, 112])
node_conv_3 = helper.make_node(
    'Conv',
    ['relu_2', 'weight_3', 'bias_3'],
    ['conv_3'],
    kernel_shape=[3,3],
    pads=[1,1,1,1],
    strides=[1,1],
)
node_relu_3 = helper.make_node(
    'Relu',
    ['conv_3'],
    ['relu_3'],
)

# layer-4   N_dim
weight_4 = helper.make_tensor('weight_4', TensorProto.FLOAT, [64, 16, 1, 1], np.random.randn(64, 16, 1, 1))
bias_4 = helper.make_tensor('bias_4', TensorProto.FLOAT, [64], np.random.randn(64))
conv_4 = helper.make_tensor_value_info('conv_4', TensorProto.FLOAT, [1, 64, 112, 112])
relu_4 = helper.make_tensor_value_info('relu_4', TensorProto.FLOAT, [1, 64, 112, 112])
node_conv_4 = helper.make_node(
    'Conv',
    ['relu_3', 'weight_4', 'bias_4'],
    ['conv_4'],
    kernel_shape=[1,1],
    pads=[0,0,0,0],
    strides=[1,1],
)
node_relu_4 = helper.make_node(
    'Relu',
    ['conv_4'],
    ['relu_4'],
)

# layer-5   H_dim
weight_5 = helper.make_tensor('weight_5', TensorProto.FLOAT, [32, 64, 3, 3], np.random.randn(32, 64, 3, 3))
bias_5 = helper.make_tensor('bias_5', TensorProto.FLOAT, [32], np.random.randn(32))
conv_5 = helper.make_tensor_value_info('conv_5', TensorProto.FLOAT, [1, 32, 112, 112])
relu_5 = helper.make_tensor_value_info('relu_5', TensorProto.FLOAT, [1, 32, 112, 112])
node_conv_5 = helper.make_node(
    'Conv',
    ['relu_4', 'weight_5', 'bias_5'],
    ['conv_5'],
    kernel_shape=[3,3],
    pads=[1,1,1,1],
    strides=[1,1],
)
node_relu_5 = helper.make_node(
    'Relu',
    ['conv_5'],
    ['relu_5'],
)

# layer-6   N_dim
weight_6 = helper.make_tensor('weight_6', TensorProto.FLOAT, [512, 32, 1, 1], np.random.randn(512, 32, 1, 1))
bias_6 = helper.make_tensor('bias_6', TensorProto.FLOAT, [512], np.random.randn(512))
conv_6 = helper.make_tensor_value_info('conv_6', TensorProto.FLOAT, [1, 512, 56, 56])
relu_6 = helper.make_tensor_value_info('relu_6', TensorProto.FLOAT, [1, 512, 56, 56])
node_conv_6 = helper.make_node(
    'Conv',
    ['relu_5', 'weight_6', 'bias_6'],
    ['conv_6'],
    kernel_shape=[1,1],
    pads=[0,0,0,0],
    strides=[2,2],
)
node_relu_6 = helper.make_node(
    'Relu',
    ['conv_6'],
    ['relu_6'],
)

# layer-7   C_dim
weight_7 = helper.make_tensor('weight_7', TensorProto.FLOAT, [512, 512, 1, 1], np.random.randn(512, 512, 1, 1))
bias_7 = helper.make_tensor('bias_7', TensorProto.FLOAT, [512], np.random.randn(512))
conv_7 = helper.make_tensor_value_info('conv_7', TensorProto.FLOAT, [1, 512, 56, 56])
relu_7 = helper.make_tensor_value_info('relu_7', TensorProto.FLOAT, [1, 512, 56, 56])
node_conv_7 = helper.make_node(
    'Conv',
    ['relu_6', 'weight_7', 'bias_7'],
    ['conv_7'],
    kernel_shape=[1,1],
    pads=[0,0,0,0],
    strides=[1,1],
)
node_relu_7 = helper.make_node(
    'Relu',
    ['conv_7'],
    ['relu_7'],
)

# layer-8   C_dim
weight_8 = helper.make_tensor('weight_8', TensorProto.FLOAT, [32, 512, 3, 3], np.random.randn(32, 512, 3, 3))
bias_8 = helper.make_tensor('bias_8', TensorProto.FLOAT, [32], np.random.randn(32))
conv_8 = helper.make_tensor_value_info('conv_8', TensorProto.FLOAT, [1, 32, 56, 56])
relu_8 = helper.make_tensor_value_info('relu_8', TensorProto.FLOAT, [1, 32, 56, 56])
node_conv_8 = helper.make_node(
    'Conv',
    ['relu_7', 'weight_8', 'bias_8'],
    ['conv_8'],
    kernel_shape=[3,3],
    pads=[1,1,1,1],
    strides=[1,1],
)
node_relu_8 = helper.make_node(
    'Relu',
    ['conv_8'],
    ['relu_8'],
)

# layer-9   N_dim
weight_9 = helper.make_tensor('weight_9', TensorProto.FLOAT, [64, 32, 3, 3], np.random.randn(64, 32, 3, 3))
bias_9 = helper.make_tensor('bias_9', TensorProto.FLOAT, [64], np.random.randn(64))
conv_9 = helper.make_tensor_value_info('conv_9', TensorProto.FLOAT, [1, 64, 28, 28])
relu_9 = helper.make_tensor_value_info('relu_9', TensorProto.FLOAT, [1, 64, 28, 28])
node_conv_9 = helper.make_node(
    'Conv',
    ['relu_8', 'weight_9', 'bias_9'],
    ['conv_9'],
    kernel_shape=[3,3],
    pads=[1,1,1,1],
    strides=[2,2],
)
node_relu_9 = helper.make_node(
    'Relu',
    ['conv_9'],
    ['relu_9'],
)

# layer-10   N_dim
weight_10 = helper.make_tensor('weight_10', TensorProto.FLOAT, [128, 64, 3, 3], np.random.randn(128, 64, 3, 3))
bias_10 = helper.make_tensor('bias_10', TensorProto.FLOAT, [128], np.random.randn(128))
conv_10 = helper.make_tensor_value_info('relu_10', TensorProto.FLOAT, [1, 128, 14, 14])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 128, 14, 14])
node_conv_10 = helper.make_node(
    'Conv',
    ['relu_9', 'weight_10', 'bias_10'],
    ['conv_10'],
    kernel_shape=[3,3],
    pads=[1,1,1,1],
    strides=[2,2],
)
node_relu_10 = helper.make_node(
    'Relu',
    ['conv_10'],
    ['output'],
)


graph_def = helper.make_graph(
    [node_conv_1, node_relu_1,
     node_conv_2, node_relu_2,
     node_conv_3, node_relu_3,
     node_conv_4, node_relu_4,
     node_conv_5, node_relu_5,
     node_conv_6, node_relu_6,
     node_conv_7, node_relu_7,
     node_conv_8, node_relu_8,
     node_conv_9, node_relu_9,
     node_conv_10, node_relu_10,],
    'test_conv_mode',
    [input],
    [output],
    initializer=[weight_1, bias_1,
                 weight_2, bias_2,
                 weight_3, bias_3,
                 weight_4, bias_4,
                 weight_5, bias_5,
                 weight_6, bias_6,
                 weight_7, bias_7,
                 weight_8, bias_8,
                 weight_9, bias_9,
                 weight_10, bias_10,],
)
mode_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.checker.check_model(mode_def)
onnx.save(mode_def, "./Conv_relu_result.onnx")

model = r'./Conv_relu_result.onnx'
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)