NVIDIA 7th SkyHackathon（七）Tao 目标检测模型可视化推理与导出

1.可视化推理结果

对于训练好的模型，使用 tao 套件可进行推理，各参数如下：

• -o：输出文件目录，即检测过的图片保存位置
• -e：训练配置文件
• -m：要测试的模型
• -l：推理结果的标注文件目录
• -k：秘钥

tao ssd inference --gpu_index=$GPU_INDEX -i $LOCAL_DATA_DIR/test_samples -o $USER_EXPERIMENT_DIR/ssd_infer_images -e $SPECS_DIR/ssd_retrain_resnet18_kitti.txt -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_resnet18_epoch_$EPOCH.tlt -l $USER_EXPERIMENT_DIR/ssd_infer_labels -k $KEY

之后，使用 matplotlib 进行可视化

import matplotlib.pyplot as plt
import os
from math import ceil
valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']

def visualize_images(image_dir, num_cols=4, num_images=10):
    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)
    num_rows = int(ceil(float(num_images) / float(num_cols)))
    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])
    f.tight_layout()
    a = [os.path.join(output_path, image) for image in os.listdir(output_path) 
         if os.path.splitext(image)[1].lower() in valid_image_ext]
    for idx, img_path in enumerate(a[:num_images]):
        col_id = idx % num_cols
        row_id = idx // num_cols
        img = plt.imread(img_path)
        axarr[row_id, col_id].imshow(img) 

效果图：

2.模型导出

首先在 $USER_EXPERIMENT_DIR 下创建模型导出的目录 export，然后使用 tao 套件导出训练的模型，导出的模型即可利用 NVIDIA TensorRT 进行推理

其参数如下：

• -m：要导出的模型
• -o：导出模型的存放位置
• -e：训练配置文件
• -k：秘钥
• –batch_size：模型的输入 batchsize
• –data_type：导出模型的数据类型

tao ssd export --gpu_index=$GPU_INDEX -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_resnet18_epoch_$EPOCH.tlt -k $KEY -o $USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt -e $SPECS_DIR/ssd_train_resnet18_kitti.txt --batch_size 1 --data_type fp32

最后，使用 tao-coverter 工具，将导出的模型转为 TensorRT 可直接使用的推理引擎即可

tao converter -k $KEY -d 3,300,300 -o NMS -e $USER_EXPERIMENT_DIR/export/trt.engine -m 1 -t fp32 -i nchw $USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt