|
下载deepsort的pytorch库,地址:https://github.com/ZQPei/deep_sort_pytorch.git
想要将reid特征提取模型 ckpt.t7转为rknn模型,大概步骤是:
1. 将ckpt.t7模型转换为onnx模型代码如下:
import argparse
import torch
import numpy as np
import onnx
from deep_sort import build_tracker
from utils.draw import draw_boxes
from utils.parser import get_config
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--config_deepsort", type=str, default="./configs/deep_sort.yaml", help='Configure tracker')
parser.add_argument("--cpu", dest="use_cuda", action="store_false", default=True, help='Run in CPU')
args = parser.parse_args()
cfg = get_config()
cfg.merge_from_file(args.config_deepsort)
use_cuda = args.use_cuda and torch.cuda.is_available()
torch.set_grad_enabled(False)
model = build_tracker(cfg, use_cuda=False)
model.reid = True
model.extractor.net.eval()
device = 'cuda' #'cuda'
output_onnx = 'ckpt.onnx'
# ------------------------ export -----------------------------
print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
input_names = ['input']
output_names = ['output']
input_tensor = torch.randn(1, 3, 128, 64, device=device)
# onnx
torch.onnx.export(model.extractor.net.cpu(), input_tensor, output_onnx, export_params=True, verbose=True,
input_names=input_names, output_names=output_names, opset_version=10,
do_constant_folding=True)
2. 验证onnx模型,结果正确。
3. 将onnx转为rknn模型,转换方法为https://github.com/EASY-EAI/yolov5 中的rknn_convert.py,
rknn正常转换成功,build成功。
4. 问题:推理ckpt.rknn模型时, rknn.inference输出的特征为负数,全部为-65536.
推理代码如下:
RKNN_MODEL = "ckpt.rknn"
IMG_PATH = "images/2_0002.jpg"
model = load_model(RKNN_MODEL)
_rknn = model
img = cv2.imread(IMG_PATH) [:, :, (2, 1, 0)]
size = (64, 128)
_im = cv2.resize(img.astype(np.float32)/255., size)
# data preprocess
ave = [0.485, 0.456, 0.406]
stda = [0.229, 0.224, 0.225]
im = Z_ScoreNorm(_im, ave, stda) # normalization
im = np.transpose(im, (2, 0, 1)) # transform 3*128*64
im = np.expand_dims(im, axis=0) # input (1*3*128*64)
im = im.astype(np.float16)
pred_onx = _rknn.inference(inputs=[im],data_type='float16', data_format='nchw') # 输入格式为nchw , data_format='nchw'
print(pred_onx)
请问有哪位大神能指点一下哪里出问题了吗? 谢谢!
reid特征提取模型,ckpt.t7
|
|