如何从0开始训练自己的CV模型

第一步配置基本环境(在上一篇已经配置了我参加第七届NVIDIA Sky Hackathon——训练ASR模型 ) 第二步利用labelimg制作图像数据集第三步开始训练resnet18模型

前言

Sky Hackathon由NVIDIA发起并主办，项目旨在帮助在校学生、深度学习开发者在NVIDIA Jetson边缘高性能计算产品上部署和优化人工智能应用。在经验丰富的GPU导师指导下，通过黑客松竞赛的方式学习业界所需的深度学习相关应用开发及其并行计算技能，激发学生们的学习兴趣与创新力。

NVIDIA工程师将亲自为参赛队伍带来他们对最新的深度学习与边缘计算方面的理解、行业的趋势与最新的技术应用及最新开发工具实战技能知识，在训练营中对参赛队伍进行指导。

Sky Hackathon为参加者提供了一个难得的学习并实操的机会，学习嵌入式深度学习开发所需的动手技能，通过使用NVIDIA最新的编程模型、库和工具以加速和优化他们的AI应用程序。

例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了机器学习的基础内容。

相关网站：第七届NVIDIA SKy Hackathon 比赛文件资料链接： https://pan.baidu.com/s/1UADeDey1OsTvEOPE5f_2CA?pwd=a647 提取码: a647

一、利用labelimg制作图像数据集

1.安装labelimg

1	pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

2.开始数据标注

在这里插入图片描述打开你的图像目录

选择Pascal VOC格式进行标注保存后会生成xml格式文件

<annotation>
	<folder>JPEGImages</folder>
	<filename>cardboard2.jpg</filename>
	<path>/your_path</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>167</width>
		<height>127</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>cardboard</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>36</xmin>
			<ymin>5</ymin>
			<xmax>131</xmax>
			<ymax>119</ymax>
		</bndbox>
	</object>
</annotation>

3.将VOC格式数据集转为KITTI数据集

# -*- coding: utf-8 -*-
import sys
from os import listdir
from os.path import isfile, isdir, join, dirname, splitext, basename
import xml.etree.ElementTree as ET

path=""
class XMLReader:
    def __init__(self, path):
        file = open(path, 'r')

        self.path = path
        self.content = file.read()
        self.root = ET.fromstring(self.content)
        self.template = "{name} 0.00 0 0.0 {xmin}.00 {ymin}.00 {xmax}.00 {ymax}.00 0.0 0.0 0.0 0.0 0.0 0.0 0.0"

    def get_filename(self):
        return splitext(basename(self.path))[0]

    def get_dir(self):
        return dirname(self.path)

    def get_objects(self):
        objects = []

        for object in self.root.findall("object"):
            objects.append({
                "name" : object.find("name").text,
                "xmin" : object.find("bndbox").find("xmin").text,
                "ymin" : object.find("bndbox").find("ymin").text,
                "xmax" : object.find("bndbox").find("xmax").text,
                "ymax" : object.find("bndbox").find("ymax").text
            })

        return objects

    def fill_template(self, object):
        return self.template.format(**object)

    def export_kitti(self):
        objects = self.get_objects()

        #Skip empty
        if len(objects) == 0: return False

        file = open(join(self.get_dir(), self.get_filename()) + ".txt", 'w')

        for object in objects[:-1]:
            file.write(self.fill_template(object) + "\n")
        # Write last without '\n'
        file.write(self.fill_template(objects[-1]))

        file.close()

        return True


def process_file(path):
    xml_reader = XMLReader(path)

    return xml_reader.export_kitti()


def get_directory_xml_files(dir):
    return [join(dir, f) for f in listdir(dir) if isfile(join(dir, f)) and splitext(f)[1].lower() == ".xml"]


def check_argv(argv):
    return len(argv) > 1


def main():
    if not check_argv(sys.argv):
        print("Wrong arguments. You should specify xml files or directory with xml files")

    # remove script name
    args = sys.argv[1:]
    processed_file_count = 0

    for path in args:
        files = []

        if isfile(path):
            files.append(path)
        elif isdir(path):
            files += get_directory_xml_files(path)

        for file in files:
            if process_file(file): processed_file_count += 1

    print("Finished. {0} Files are processed".format(processed_file_count))

if __name__ == "__main__":
    main()

二、开始训练resnet18模型

1. 设置训练环境和路径

代码如下（示例）：当我们使用训练模型的时候，需要从NGC官网申请一个API Key。这个Key将和您训练的模型进行绑定。 Tao Toolkit可以直接在您计算机系统的环境下执行，它会自动启动docker，并加载相关工具。我们需要设置两套文件地址：
LOCAL_PROJECT_DIR-->您计算机系统中的实验文档存放地址
USER_EXPERIMENT_DIR-->docker系统中的实验文档存放地址

在下面的步骤中，我们会将他们形成映射，您只需要操作你计算机系统中地址就好，无需管理docker系统中的地址

# Setting up env variables for cleaner command line commands.
import os

print("Please replace the variable with your key.")
%env KEY=your_key
%env GPU_INDEX=0
%env USER_EXPERIMENT_DIR=/your_path/CV
%env DATA_DOWNLOAD_DIR=/your_path/CV/data

%env LOCAL_PROJECT_DIR=/your_path/CV
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()))

os.environ["LOCAL_SPECS_DIR"] = os.path.join(
    os.getenv("NOTEBOOK_ROOT", os.getcwd()),
    "specs"
)
%env SPECS_DIR=/your_path/7th/specs

!echo $LOCAL_SPECS_DIR
!ls -rlt $LOCAL_SPECS_DIR

2.将系统目录映射到Tao Docker

代码如下（示例）：

# 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": "/home/meng/7thSkyHackathon/CV/"
        },
        # Mapping the specs directory.
        {
            "source": os.environ["LOCAL_SPECS_DIR"],
            "destination": os.environ["SPECS_DIR"]
        },
    ]
}

# Writing the mounts file.
with open(mounts_file, "w") as mfile:
    json.dump(drive_map, mfile, indent=4)

查看目录下文件

1	!cat ~/.tao_mounts.json

3.安装TAO的加载器

1
2
3

# 如果您已经安装好了，则不需要执行这一步
!pip3 install nvidia-pyindex
!pip3 install nvidia-tao

查看TAO的信息

1 2	# 查看当前环境下TAO的信息 !tao info

## 4.准备数据集和预训练模型

# verify
# 此时的LOCAL_DATA_DIR为：LOCAL_PROJECT_DIR/data
# 比如我的LOCAL_PROJECT_DIR为/home/hekun/mydata/workspace/tao-experiments， 那么在我的实验文档中LOCAL_DATA_DIR=/home/hekun/mydata/workspace/tao-experiments/data
# 您需要将您的图片数据集（此时我的为pascal_img）和标注文件（此时我的为pascal_label）放到LOCAL_DATA_DIR中
!ls -l $LOCAL_DATA_DIR/

# 生成验证数据集，主要您要把generate_val_dataset.py文件中图片数据的格式后缀名修改成跟您图片数据格式一样。原文中使用的是.png格式的后缀名。
# 此时generate_val_dataset.py的参数需要的文件夹地址均为您的计算机系统本地地址，不是docker中的地址。
!python3.8 generate_val_dataset.py --input_image_dir=$LOCAL_DATA_DIR/images \
                                   --input_label_dir=$LOCAL_DATA_DIR/labels \
                                   --output_dir=$LOCAL_DATA_DIR/val

#修改specs/ssd_tfrecords_kitti_train.txt中的内容：
#1.root_directory_path为docker环境中的地址，你需要修改和你自己的系统环境相对应。比如在我的例子中，
#  保存图片数据的pascal_img文件夹在/home/hekun/mydata/workspace/tao-experiments/data目录下。
#  而/home/hekun/mydata/workspace/tao-experiments/正好映射到docker文件系统目录中的/workspace/tao-experiments/。 所以我就将root_directory_path-->/workspace/tao-experiments/data
#2.image_dir_name定义为保存图片数据的文件夹，您可以理解为Tao会在root_directory_path的目录下寻找image_dir_name文件夹，image_dir_name文件夹里面的所有图片都被视作训练图片
#3.label_dir_name定义为保存标注文件的文件夹。
#4.image_directory_path 和root_directory_path可以理解为一样。
#5.target_class_mapping定义为您标注的类别，key为您标注文件中写的类别，value为Tao训练时使用的类别名字
print("TFRecords conversion spec file:")
!cat $LOCAL_SPECS_DIR/ssd_tfrecords_kitti_train.txt

# 创建用于保存tfrecord格式的数据的文件夹，并利用Tao自带的工具转换数据。数据将被保存在-o后面的文件目录里面。
print("Converting the training set to TFRecords.")
!mkdir -p $LOCAL_DATA_DIR/tfrecords && rm -rf $LOCAL_DATA_DIR/tfrecords/*
!tao ssd dataset_convert \
         -d /your_path/CV/7th/specs/ssd_tfrecords_kitti_train.txt \
         -o /your_path/CV/data/tfrecords/kitti_train

查看目录下文件

1	!ls -rlt $LOCAL_DATA_DIR/tfrecords/

下载预训练模型

#下面的步骤是用来下载并安装NGC CLI的，第一次运行的时候需要安装，如果已经安装了，那么可以跳过这一步
# 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", ""))

1
2

#此步骤是用来验证NGC的，如果下载不成功出现问题，可能是您没有开启NGC的验证，或者您的Key过时了，需要重新验证。如果下载成功，可以不必运行这一步。
!docker login nvcr.io

1 2	#同上一步，如果下面一步下载不成功，可以清理之前的验证，并运行上一步重新验证。如果下载成功，可以不必运行这一步。 !ngc config clear

1 2	#查看目前NGC上可以下载的预训练模型，比如下面这个命令就是列出所有的目标检测的预训练模型 !ngc registry model list nvidia/tao/pretrained_object_detection:*

1 2	#创建文件夹保存下载的预训练模型 !mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/

1 2	# 下载您所选择的预训练模型 !ngc registry model download-version nvidia/tao/pretrained_object_detection:resnet18 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18

1
2
3

#查看预训练模型是否下载成功
print("Check that model is downloaded into dir.")
!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_object_detection_vresnet18

5.定义训练的参数

修改target_class_mapping为您需要训练的类别
此处的路径基本不需要修改，除非您自己改动过。
batch_size_per_gpu定义了batch_size大小，如果您的GPU显存不是很大，建议可以调整小一点，否则可能出现out of memory的错误
num_epochs定义了您会训练多少圈，如果您第一次训练，建议不要小于80
validation_period_during_training定义了训练多少圈验证一下，可以直观地看出您训练的效果变化，可以定义5 or 10

通过修改此处的txt来修改训练的参数

1	!cat $LOCAL_SPECS_DIR/ssd_train_resnet18_kitti.txt

## 6.开始利用TAO训练

您只需要利用下面命令中的tao ssd train就可以训练
--gpus定义了您有几个GPU，--gpu_index定义了您用哪个GPU
-e定义了您上面设置的训练设置文件
-r定义了您训练的模型输出文件夹
-k定义了您的秘钥API Key
-m定义了您的预训练模型

1 2	#创建模型输出文件夹，您训练的结果会保存在这里 !mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned_final

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_final \
               -k $KEY \
               -m $USER_EXPERIMENT_DIR/7th/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_020.tlt \
                --initial_epoch 20

#查看您训练的模型
print('Model for each epoch:')
print('---------------------')
!ls -ltrh $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/weights

找到效果最好的那个模型

#查看您训练模型的精度，并选择一个最好的，把那个最好的模型的epoch复制个下面的EPOCH，说明我们后面的操作都是基于这个模型。
#这个结果中，倒数第二个参数为mAP值，即您训练模型的精度，最后一个值为验证的损失。我们会选择mAP最大的那个模型
!cat $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/ssd_training_log_resnet18.csv
%set_env EPOCH=000

7.评估训练的模型

#评估训练的模型
!tao ssd evaluate --gpu_index=$GPU_INDEX \
                  -e $SPECS_DIR/ssd_train_resnet18_kitti.txt \
                  -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/weights/ssd_resnet18_epoch_$EPOCH.tlt \
                  -k $KEY

8.对模型进行剪枝

-m 定义您要剪枝的模型
-o 定义剪之后输出的模型
-eq定义剪枝的方式
-pth定义剪枝的阈值。这里可以理解为，pth越大剪去参数越多，模型越小，运行速度越快，但是可能精度越低。所以采用多大的pth需要反复试验，建议从0.5开始，如果你的模型训练的效果一直不好，那么把0.5换成0.3，减少剪枝的部分。
-k 定义您的秘钥API Key

1 2	%set_env EPOCH=020 !mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_pruned

pth是你剪枝的参数，pth大则模型小，精度不高，pth大则模型精度高，文件大

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

1	!ls -rlt $USER_EXPERIMENT_DIR/experiment_dir_pruned/

9.对剪枝后的模型重新训练

这里就像刚才训练时一样，我们需要定义训练的设置
注意这里的训练设置文件和刚才的训练设置文件不一样
您还是需要像上面的步骤一样修改其中的一些参数

1 2	#查看重新训练的设置文件 !cat $LOCAL_SPECS_DIR/ssd_retrain_resnet18_kitti.txt

1 2	#创建文件夹，用于保存训练结果 !mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_retrain

# 跟上面训练的步骤一样，只不过修改了预训练模型
!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

1 2	# 查看训练结果 !ls -rlt $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights

1
2
3

# 同样的步骤，找到那个最好的
!cat $USER_EXPERIMENT_DIR/experiment_dir_retrain/ssd_training_log_resnet18.csv
%set_env EPOCH=000

10.评估新的模型

!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

11.可视化推理结果

1	!mkdir -p $LOCAL_DATA_DIR/test_samples

1	!ls $LOCAL_DATA_DIR

1 2	!echo $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_resnet18_epoch_$EPOCH.tlt !ls $LOCAL_DATA_DIR/test_samples

# 在此步骤中：
#1.我们需要将要测试的图片放在-i后面的文件夹中，注意这里的文件夹是Tao docker文件系统中的路径
#2.-o是输出文件夹，即检测过得图片保存的位置
#3.-e是您训练时用的文件
#4.-m是您要测试的模型
#5.-l是推理结果的标注文件
#6.-k是您的秘钥API Key
!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

1	!pip3 --default-timeout=100 install matplotlib==3.3.3

# 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_cols
        row_id = idx // num_cols
        img = plt.imread(img_path)
        axarr[row_id, col_id].imshow(img)

1	!ls $USER_EXPERIMENT_DIR/ssd_infer_images

# Visualizing the sample images.
OUTPUT_PATH = '/home/meng/7thSkyHackathon/CV/ssd_infer_images' # relative path from $USER_EXPERIMENT_DIR.
COLS = 3 # number of columns in the visualizer grid.
IMAGES = 21 # number of images to visualize.

visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)

12.模型的导出

通过tao工具导出您训练的模型，导出的模型可以利用NVIDIA TensorRT或者NVIDIA Deepstream进行推理部署：
-m定义了您要导出的模型
-k定义了您的秘钥API Key
-o 定义了您导出的模型存放的位置
-e定义了您训练这个模型的设置文件
--batch_size定义了模型的输入batchsize
--data_type定义了您导出的模型的数据类型，可以定义为fp32，fp16 or int8。如果定义了int8 您还需要定义下面的量化文件。

!rm -rf $USER_EXPERIMENT_DIR/export
!mkdir -p $USER_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_train_resnet18_kitti.txt \
                --batch_size 1 \
                --data_type fp32

1
2
3

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

在这里利用tao-converter工具将您训练的模型转化为TensorRT可以直接使用的推理引擎，注意：您需要在哪里使用您的模型，您就要在哪里转换这个模型。
您可以访问下面的地址，下载tao-converter工具：

注意：本地导出的模型不可用在节点上，两者代码虽然一样，但是创建的模型是不一样的

!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

1
2
3

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

13.验证导出的模型

!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

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)

总结

在训练过程中遇到各种各样的问题要善于去自己改正，有些很简单的问题可能只是自己疏忽大意导致的，一点一点debug才能进步

我参加第七届NVIDIA Sky Hackathon——训练CV模型