# ---------------------------------------------------------------------#
# 定义一个简单的单输入网络，地平线要求模型中必须要有Conv/Gemm/MatMul中的一个
# ---------------------------------------------------------------------#
class MyNet(nn.Module):
def __init__(self, num_classes=10):
super(MyNet, self).init()
self.fc = nn.Linear(162, num_classes)

# -----------------------------------#
# 导出ONNX模型函数
# -----------------------------------#
def model_convert_onnx(model, output_path):
a = torch.randn(1, 3, 3, 3, 3)
b = torch.randn(1, 3, 3, 3, 3)
input_names = ["input1", "input2"] # 导出的ONNX模型输入节点名称
output_names = ["output1"] # 导出的ONNX模型输出节点名称

if name == '__main__':
model = MyNet()
# print(model)

# 模型输入相关参数, 若输入多个节点, 则应使用';'进行分隔, 使用默认缺省设置则写None
input_parameters:
# 模型输入的节点名称, 此名称应与模型文件中的名称一致, 否则会报错, 不填则会使用模型文件中的节点名称
input_name: "input1;input2"
# 网络实际执行时，输入给网络的数据格式，包括 nv12/rgb/bgr/yuv444/gray/featuremap
input_type_rt: 'featuremap;featuremap'
# 网络实际执行时输入的数据排布, 可选值为 NHWC/NCHW
input_layout_rt: 'NHWC;NHWC'
# 网络训练时输入的数据格式，可选的值为rgb/bgr/gray/featuremap/yuv444
input_type_train: 'featuremap;featuremap'
# 网络训练时输入的数据排布, 可选值为 NHWC/NCHW
input_layout_train: 'NHWC;NHWC'
# 网络输入的预处理方法，主要有以下几种：
# no_preprocess 不做任何操作
norm_type: 'no_preprocess;no_preprocess'

# 模型量化相关参数
calibration_parameters:
calibration_type: 'skip' # 不涉及校准数据

# 编译器相关参数
compiler_parameters:
# 编译策略，支持bandwidth和latency两种优化模式
compile_mode: 'latency'
# 优化等级可选范围为O0~O3
optimize_level: 'O0'

image.png