Toybrick

标题: 关于Slice的问题如何解决 [打印本页]

作者: phker 时间: 2021-11-2 10:30
标题: 关于Slice的问题如何解决
.[attach]2119[/attach]

模型问题太大无法上传
源代码如下
#!./bin/python3.6
# -*- coding:utf-8 -*-

from rknn.api import RKNN

INPUT_SIZE = 64

if __name__ == '__main__':
# 创建RKNN执行对象
rknn = RKNN()
# 配置模型输入，用于NPU对数据输入的预处理
# channel_mean_value='0 0 0 255'，那么模型推理时，将会对RGB数据做如下转换
# (R - 0)/255, (G - 0)/255, (B - 0)/255。推理时，RKNN模型会自动做均值和归一化处理
# reorder_channel=’0 1 2’用于指定是否调整图像通道顺序，设置成0 1 2即按输入的图像通道顺序不做调整
# reorder_channel=’2 1 0’表示交换0和2通道，如果输入是RGB，将会被调整为BGR。如果是BGR将会被调整为RGB
#图像通道顺序不做调整
# rknn.config(channel_mean_value='0 0 0 255', reorder_channel='0 1 2')
# rknn.config(
#    mean_values=[[123.675, 116.28, 103.53]],
#    std_values=[[58.395, 58.395, 58.395]],
#    reorder_channel='0 1 2')
add_perm = False # 如果设置成True,则将模型输入layout修改成NHWC
rknn.config(
      batch_size=1,
      mean_values=[[0, 0, 0]],
      std_values=[[255, 255, 255]],
      reorder_channel='0 1 2',
      force_builtin_perm=add_perm,
      output_optimize=1,
      target_platform='rk3399pro')

# 加载TensorFlow模型
# tf_pb='digital_gesture.pb'指定待转换的TensorFlow模型
# inputs指定模型中的输入节点
# outputs指定模型中输出节点
# input_size_list指定模型输入的大小
print('--> Loading model')
# rknn.load_tensorflow(tf_pb='digital_gesture.pb',
#                   inputs=['input_x'],
#                   outputs=['probability'],
#                   input_size_list=[[INPUT_SIZE, INPUT_SIZE, 3]])

# rknn.load_onnx("/home/roota/Desktop/AI/rknntools/yolov5s.onnx")
# print('done')
ret = rknn.load_onnx(
      model='yolov5m_epoch10sim.onnx',
      inputs=['images'],
      input_size_list=[[3,640,640]],
      outputs=['output'])
if ret != 0:
      print('Load  onnx  failed!')
      exit(ret)
print('done')

# 创建解析pb模型
# do_quantization=False指定不进行量化
# 量化会减小模型的体积和提升运算速度，但是会有精度的丢失
print('--> Building model')
rknn.build(do_quantization=True)
print('done')

# 导出保存rknn模型文件
ret = rknn.export_rknn('./yolov5m_epoch10.rknn')
if ret != 0:
      print('Export yolov5s.rknn failed!')
      exit(ret)
else:
      print('Export yolov5s.rknn success!')

# Release RKNN Context
rknn.release()

欢迎光临 Toybrick (https://t.rock-chips.com/)