车牌检测模型训练(含源码和数据集)

NVIDIA TAO是以docker镜像的形式运行, 我们需要设置自己的计算机系统和docker镜像内系统的路径映射关系.

我们采用了clearml的可视化技术, 可以利用clearml来监测我们的训练过程, 这里还需要设置clearml相关的配置信息

import osprint("Please replace the variable with your key.")
#设置NVIDIA TAO的api key, 这个key是用来加密你的模型的
%env KEY=nvidia_tlt
%env GPU_INDEX=0#设置实验路径和数据集路径, 这里的路径指的是docker镜像内部的路径
%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/ssd
%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data#设置clearml的相关配置信息
%env CLEARML_WEB_HOST=https://app.clear.ml
%env CLEARML_API_HOST=https://api.clear.ml
%env CLEARML_FILES_HOST=https://files.clear.ml
%env CLEARML_API_ACCESS_KEY=SY1OMSO6JY6W24WVBHOR
%env CLEARML_API_SECRET_KEY=URcjcox7nx9YIE7mDFoabr1CsdVcMWKSU0NPb5machH3aPP8FZ#这里设置的路径是指的你自己计算机系统的路径, 当前设置的是我的, 你需要根据你自己的实验路径更改
%env LOCAL_PROJECT_DIR=/home/hekun/mydata/tao-3.0-sample/cv_samples_v1.4.0/ssd
os.environ["LOCAL_DATA_DIR"] = os.path.join(os.getenv("LOCAL_PROJECT_DIR", os.getcwd()), "data")
os.environ["LOCAL_EXPERIMENT_DIR"] = os.path.join(os.getenv("LOCAL_PROJECT_DIR", os.getcwd()), "ssd")#这里设置的是训练,数据转换等定义文件的存放路径
os.environ["LOCAL_SPECS_DIR"] = os.path.join(os.getenv("NOTEBOOK_ROOT", os.getcwd()),"specs"
)
%env SPECS_DIR=/workspace/tao-experiments/ssd/specs!ls -rlt $LOCAL_SPECS_DIR

下面的单元格将本地主机上的项目目录映射到 TAO docker 实例中的工作区目录，以便数据和结果从 docker 中映射出来。 更多信息请参考用户指南中的启动器实例。

envs里面指的是clearml可视化的设置信息, 你需要根据你自己的账户设置CLEARML_API_ACCESS_KEY和CLEARML_API_SECRET_KEY

# Mapping up the local directories to the TAO docker.
import json
mounts_file = os.path.expanduser("~/.tao_mounts.json")# Define the dictionary with the mapped drives
drive_map = {"Mounts": [# Mapping the data directory{"source": os.environ["LOCAL_PROJECT_DIR"],"destination": "/workspace/tao-experiments"},# Mapping the specs directory.{"source": os.environ["LOCAL_SPECS_DIR"],"destination": os.environ["SPECS_DIR"]},],"Envs": [{"variable": "CLEARML_WEB_HOST","value": "https://app.clear.ml"},{"variable": "CLEARML_API_HOST","value": "https://api.clear.ml"},{"variable": "CLEARML_FILES_HOST","value": "https://files.clear.ml"},{"variable": "CLEARML_API_ACCESS_KEY","value": "SY1OMSO6JY6W24WVBHOR"},{"variable": "CLEARML_API_SECRET_KEY","value": "URcjcox7nx9YIE7mDFoabr1CsdVcMWKSU0NPb5machH3aPP8FZ"}],"DockerOptions": {"shm_size": "16G","ulimits": {"memlock": -1,"stack": 67108864},"user": "1000:1000","ports": {"8888": 8888}}
}# Writing the mounts file.
with open(mounts_file, "w") as mfile:json.dump(drive_map, mfile, indent=4)

!cat ~/.tao_mounts.json

TAO 启动器是一个 python 包，作为 PyPI 中列出的 python wheel 分发。 您可以通过执行以下单元格来安装启动器。

请注意，TAO Toolkit 建议用户使用 python 3.6.9 在虚拟环境中运行 TAO 启动器。 您可以按照此 页面 中的说明使用 virtualenv 和 virtualenvwrapper 包设置 python 虚拟环境。 设置 virtualenvwrapper 后，请使用 VIRTUALENVWRAPPER_PYTHON 变量设置要在虚拟环境中使用的 python 版本。 你可以通过运行来做到这一点

export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x

其中 x >= 6 且 <= 8

我们建议先执行此步骤，然后从虚拟环境启动笔记本。 除了安装 TAO python 包外，请确保满足以下软件要求：

• python >=3.6.9 < 3.8.x
• docker-ce > 19.03.5
• docker-API 1.40
• nvidia-container-toolkit > 1.3.0-1
• nvidia-container-runtime > 3.4.0-1
• nvidia-docker2 > 2.5.0-1
• nvidia-driver > 455+

安装先决条件后，请按照以下命令登录到 docker registry nvcr.io

docker login nvcr.io

您将被触发输入用户名和密码。 用户名是“$oauthtoken”，密码是从“ngc.nvidia.com”生成的 API 密钥。 请按照 NGC 设置指南 中的说明生成您自己的 API 密钥。

# SKIP this step IF you have already installed the TAO launcher.
!pip3 install nvidia-tao

# View the versions of the TAO launcher
!tao info

CRPD (Chinese Road Plate Dataset)是一个大型的中文车牌数据集, 它包含三个子集— CRPD-single, CRPD-double ,CRPD-multi. 本教程采用了其中的CRPD-single为例. 他的标注包含:

x1, y1, x2, y2, x3, y3, x4, y4, type, content

• 前八个数字表示边界四边形角的坐标。
• “type”标注表示LP的类型，0代表蓝盘，1代表黄盘和单行车牌，2代表黄牌和双行车牌，3代表白牌。
• “content”注释表示 LP 内容。

由于NVIDIA TAO目标检测模型训练数据集需要KITTI格式, 所以我们需要将数据集的标注进行转换.我已经将转换格式的脚本放在本教程文件夹中, 请查看preprocess_crpd.py文件, 并根据你实际的路径作出简单调整就可以使用

!python3 preprocess_crpd.py

接下来将准备好的数据集分割一部分出来当做测试集

# Generate val dataset out of training dataset
!python3 generate_val_dataset.py --input_image_dir=$LOCAL_DATA_DIR/lpd_images \--input_label_dir=$LOCAL_DATA_DIR/lpd_labels \--output_dir=$LOCAL_DATA_DIR/val

当我们准备好数据集之后, 我们需要将数据集转换成tfrecords格式, 这里可以利用tao自带的工具.

我们在这里需要定义数据的路径, 格式, 目标类别等信息.

需要注意的是这里的路径root_directory_path指的是docker系统内的路径

print("TFRecords conversion spec file:")
!cat $LOCAL_SPECS_DIR/ssd_tfrecords_kitti_train.txt

开始转换

# Creating a new directory for the output tfrecords dump.
print("Converting the training set to TFRecords.")
!mkdir -p $LOCAL_DATA_DIR/tfrecords && rm -rf $LOCAL_DATA_DIR/tfrecords/*
!tao ssd dataset_convert \-d $SPECS_DIR/ssd_tfrecords_kitti_train.txt \-o $DATA_DOWNLOAD_DIR/tfrecords/kitti_train
!ls -rlt $LOCAL_DATA_DIR/tfrecords/

我们将使用 NGC CLI 获取预训练模型。 有关详细信息，请访问 ngc.nvidia.com 并单击导航栏上的 SETUP。

这里如果因为网络的原因没法下载, 也可以手动登录并下载模型:

https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tao/models/pretrained_object_detection/files

# Installing NGC CLI on the local machine.
## Download and install
%env CLI=ngccli_cat_linux.zip
!mkdir -p $LOCAL_PROJECT_DIR/ngccli# Remove any previously existing CLI installations
!rm -rf $LOCAL_PROJECT_DIR/ngccli/*
!wget "https://ngc.nvidia.com/downloads/$CLI" -P $LOCAL_PROJECT_DIR/ngccli
!unzip -u "$LOCAL_PROJECT_DIR/ngccli/$CLI" -d $LOCAL_PROJECT_DIR/ngccli/
!rm $LOCAL_PROJECT_DIR/ngccli/*.zip 
os.environ["PATH"]="{}/ngccli:{}".format(os.getenv("LOCAL_PROJECT_DIR", ""), os.getenv("PATH", ""))

!ngc registry model list nvidia/tao/pretrained_object_detection:*

!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/

# Pull pretrained model from NGC
!ngc registry model download-version nvidia/tao/pretrained_object_detection:resnet18 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18

print("Check that model is downloaded into dir.")
!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_object_detection_vresnet18

我们需要定义训练的一些参数, 通过scpecs/ssd_train_resnet18_kitti.txt文件来定义:

!cat $LOCAL_SPECS_DIR/ssd_train_resnet18_kitti.txt

接下来我们开始训练, 这个过程会比较久, 你也可以通过clearml来监控训练过程

!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned

print("To run with multigpu, please change --gpus based on the number of available GPUs in your machine.")
!tao ssd train --gpus 1 --gpu_index=$GPU_INDEX \-e $SPECS_DIR/ssd_train_resnet18_kitti.txt \-r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \-k $KEY \-m $USER_EXPERIMENT_DIR/pretrained_resnet18/pretrained_object_detection_vresnet18/resnet_18.hdf5

print("To resume from checkpoint, please uncomment and run this instead. Change last two arguments accordingly.")
# !tao ssd train --gpus 1 --gpu_index=$GPU_INDEX \
#                -e $SPECS_DIR/ssd_train_resnet18_kitti.txt \
#                -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \
#                -k $KEY \
#                -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/ssd_resnet18_epoch_001.tlt \
#                --initial_epoch 2

print('Model for each epoch:')
print('---------------------')
!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/weights

# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.
!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/ssd_training_log_resnet18.csv
%set_env EPOCH=080

!tao ssd evaluate --gpu_index=$GPU_INDEX \-e $SPECS_DIR/ssd_train_resnet18_kitti.txt \-m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/ssd_resnet18_epoch_$EPOCH.tlt \-k $KEY

• 指定预训练模型
• 设置剪枝阈值

通常，您只需要调整 -pth（阈值）以获得准确性和模型大小的权衡。 更高的 pth 给你更小的模型（因此推理速度更快）但准确性更差。 阈值取决于数据集和模型。 下面块中的“0.5”只是一个起点。 如果重新训练的准确性很好，您可以增加此值以获得更小的模型。 否则，降低此值以获得更好的准确性。

!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned

!tao ssd prune --gpu_index=$GPU_INDEX \-m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/ssd_resnet18_epoch_$EPOCH.tlt \-o $USER_EXPERIMENT_DIR/experiment_dir_pruned/ssd_resnet18_pruned.tlt \-eq intersection \-pth 0.1 \-k $KEY

!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned/

• 剪枝后需要重新训练模型以恢复准确率
• 指定再训练参数
• 警告：训练需要数小时或一天才能完成

# Printing the retrain spec file. 
# Here we have updated the spec file to include the newly pruned model as a pretrained weights.
!cat $LOCAL_SPECS_DIR/ssd_retrain_resnet18_kitti.txt

!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain
# Retraining using the pruned model as pretrained weights 
!tao ssd train --gpus 1 --gpu_index=$GPU_INDEX \-e $SPECS_DIR/ssd_retrain_resnet18_kitti.txt \-r $USER_EXPERIMENT_DIR/experiment_dir_retrain \-m $USER_EXPERIMENT_DIR/experiment_dir_pruned/ssd_resnet18_pruned.tlt \-k $KEY

# Listing the newly retrained model.
!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/weights
# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.
!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/ssd_training_log_resnet18.csv
%set_env EPOCH=080

!tao ssd evaluate --gpu_index=$GPU_INDEX \-e $SPECS_DIR/ssd_retrain_resnet18_kitti.txt \-m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_resnet18_epoch_$EPOCH.tlt \-k $KEY

在本节中，我们运行 infer 工具来生成对训练模型的推理并可视化结果。

# Running inference for detection on n images
!tao ssd inference --gpu_index=$GPU_INDEX -i $DATA_DOWNLOAD_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

tao 推理工具产生两个输出。

1. 在 $USER_EXPERIMENT_DIR/ssd_infer_images 中叠加图像
2. kitti 格式的逐帧 bbox 标签位于 $USER_EXPERIMENT_DIR/ssd_infer_labels

# Simple grid visualizer
!pip3 install matplotlib==3.3.3
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_colsrow_id = idx // num_colsimg = plt.imread(img_path)axarr[row_id, col_id].imshow(img) # Visualizing the sample images.
OUTPUT_PATH = 'ssd_infer_images' # relative path from $USER_EXPERIMENT_DIR.
COLS = 3 # number of columns in the visualizer grid.
IMAGES = 9 # number of images to visualize.visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)

如果您训练了非 QAT 模型，您可以使用下面的代码块以 FP32、FP16 或 INT8 模式导出。 对于 INT8，您需要提供校准图像目录。

# tao <task> export will fail if .etlt already exists. So we clear the export folder before tao <task> export
!rm -rf $LOCAL_EXPERIMENT_DIR/export
!mkdir -p $LOCAL_EXPERIMENT_DIR/export
# Export in FP32 mode. Change --data_type to fp16 for FP16 mode
!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_retrain_resnet18_kitti.txt \--batch_size 16 \--data_type fp32 \--gen_ds_config# Uncomment to export in INT8 mode (generate calibration cache file).
# !tao ssd export --gpu_index=$GPU_INDEX \
#                 -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_resnet18_epoch_$EPOCH.tlt  \
#                 -o $USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt \
#                 -e $SPECS_DIR/ssd_retrain_resnet18_kitti.txt \
#                 -k $KEY \
#                 --cal_image_dir $DATA_DOWNLOAD_DIR/testing/image_2 \
#                 --data_type int8 \
#                 --batch_size 16 \
#                 --batches 10 \
#                 --cal_cache_file $USER_EXPERIMENT_DIR/export/cal.bin  \
#                 --cal_data_file $USER_EXPERIMENT_DIR/export/cal.tensorfile \
#                 --gen_ds_config

print('Exported model:')
print('------------')
!ls -lh $LOCAL_EXPERIMENT_DIR/export

使用 docker 附带的“tao-converter”实用程序验证引擎生成。

tao-converter 为它所在的平台生成优化的 tensorrt 引擎。 因此，为了获得最佳性能，请实例化此 docker 并在目标设备上使用导出的 .etlt 文件和校准缓存（用于 int8 模式）执行 tao-converter 命令。 此 docker 中包含的 tao-converter 实用程序仅适用于具有独立 NVIDIA GPU 的 x86 设备。

对于 jetson 设备，请从开发区链接 此处 下载 tao-converter for jetson。

如果您选择将您的模型直接集成到 deepstream 中，您可以通过简单地将导出的 .etlt 文件连同校准缓存复制到目标设备并更新配置 gst-nvinfer 元素指向的 spec 文件来实现 到这个新导出的模型。 通常，此文件对于检测模型称为“config_infer_primary.txt”，对于分类模型称为“config_infer_secondary_*.txt”。

# Convert to TensorRT engine (FP32)
!tao converter -k $KEY \-d 3,300,300 \-o NMS \-e $USER_EXPERIMENT_DIR/export/trt.engine \-m 16 \-t fp32 \-i nchw \$USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt# Convert to TensorRT engine (FP16)
# !tao converter -k $KEY \
#                    -d 3,300,300 \
#                    -o NMS \
#                    -e $USER_EXPERIMENT_DIR/export/trt.engine \
#                    -m 16 \
#                    -t fp16 \
#                    -i nchw \
#                    $USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt# Convert to TensorRT engine (INT8).
# !tao converter -k $KEY  \
#                    -d 3,300,300 \
#                    -o NMS \
#                    -c $USER_EXPERIMENT_DIR/export/cal.bin \
#                    -e $USER_EXPERIMENT_DIR/export/trt.engine \
#                    -b 8 \
#                    -m 16 \
#                    -t int8 \
#                    -i nchw \
#                    $USER_EXPERIMENT_DIR/export/ssd_resnet18_epoch_$EPOCH.etlt

print('Exported engine:')
print('------------')
!ls -lh $LOCAL_EXPERIMENT_DIR/export/trt.engine

通过可视化 TensorRT 推理来验证转换后的引擎。

# Infer using TensorRT engine# The engine batch size once created, cannot be alterred. So if you wish to run with a different batch-size,
# please re-run tlt-convert.!tao ssd inference --gpu_index=$GPU_INDEX \-m $USER_EXPERIMENT_DIR/export/trt.engine \-e $SPECS_DIR/ssd_retrain_resnet18_kitti.txt \-i $DATA_DOWNLOAD_DIR/test_samples \-o $USER_EXPERIMENT_DIR/ssd_infer_images \-t 0.4

# Visualizing the sample images.
OUTPUT_PATH = 'ssd_infer_images' # relative path from $USER_EXPERIMENT_DIR.
COLS = 3 # number of columns in the visualizer grid.
IMAGES = 9 # number of images to visualize.visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)