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