0 背景

NVIDIA Transfer Learning Toolkit是英伟达推出的迁移学习工具包，使用提供的预训练模型来训练自己的数据，并可以完成模型裁剪、导出到deepstream框架中，实现快速部署。流程如下：

归纳如下：

下载预训练模型

数据集准备，转化为KITTI数据集格式

使用 tlt-train 工具训练

使用 tlt-evaluate 和 tlt-infer 工具验证模型精度

使用 tlt-prune 工具裁剪模型，通过减少模型参数来提升模型 FPS

使用 tlt-train 工具重新训练

使用 tlt-evaluate 工具验证 re-trained 的模型

如果模型精度满足不了需求，重新执行第5/6/7个步骤，否则到第9步

使用 tlt-export 工具导出模型，根据平台选择不同精度，jetson端只有xavier支持int8

本文以训练faster rcnn为例，对TLT的使用方法进行介绍，并记录我在部署过程中踩过的那些（深）坑，如果你要训练其它模型，方法是类似的。参考资源

官网介绍：https://developer.nvidia.com/transfer-learning-toolkit

开发文档：https://docs.nvidia.com/metropolis/TLT/tlt-getting-started-guide/index.html

NGC官网：https://ngc.nvidia.com/catalog/collections

版本信息：v2.0_dp_py2，更新时间：2020年5月7日

注：2020.08.05 v2.0_py3版本更新，使用方法类似

1 安装方法

1.1 拉取镜像

TLT需要在docker中运行，因此第一步拉取镜像

首先登陆NGC，如果你是第一次注册NGC，你需要获得一个API KEY，这个值需要保存下来，因为他只会显示一次，并且以后会经常用到

docker login nvcr.io
# 用户名固定为  $oauthtoken
# 密码为你自己的API_KEY
docker pull nvcr.io/nvidia/tlt-streamanalytics:v2.0_dp_py2

这个过程比较漫长，第一次拉完之后，如果有自己的镜像仓库，可以推送到本地镜像仓库，以后下载时就会快很多

1.2 启动容器

镜像下载好之后，启动容器

sudo nvidia-docker run --gpus 4 -it --name tlt -v /your/local/path/tlt-experiments/:/workspace/tlt-experiments -p 8888:8888 nvcr.io/nvidia/tlt-streamanalytics:v2.0_dp_py2:v2.0_dp_py2 /bin/bash

有几个点需要注意：

--gpus要根据自己硬件条件设置，我这里有4张显卡可以用，就设置为4
-v表示将自己本地目录/your/local/path/tlt-experiments/与容器中的/workspace/tlt-experiments目录绑定，这样方便拷贝我们的训练数据
-p 8888:8888表示将容器内的8888端口与外边物理机的端口进行绑定，如果8888端口已经被占用，可以修改为其它端口，如-p 8889:8888，这样启动jupyter notebook的时候也要指定8889端口
如果要指定使用哪一块GPU，使用-e NVIDIA_VISIBLE_DEVICES=4，5环境变量来设置

如果已经创建过容器，首先使用 docker ps -a 找到容器id，然后依次使用 docker start 和 docker attach 命令进入容器

1.3 修改entrypoint.sh文件

2.0DP版本的镜像有个bug，就是当你启动容器退出后无法再次进入，根据docker logs tlt可以看到报错信息是：

mkdir: cannot create directory ‘/opt/ngccli’: File exists（具体看参考论坛提问）

并且会重复下载ngccli_reg_linux.zip文件，这里的解决方法是注释掉相关代码，entrypoint.sh文件在进入容器后的上一级目录（cd ..）

#!/usr/bin/env bash
set -e## Run startup command
#mkdir -p /opt/ngccli
#wget https://ngc.nvidia.com/downloads/ngccli_reg_linux.zip -P /opt/ngccli
#unzip -o /opt/ngccli/ngccli_reg_linux.zip -d /opt/ngccli/
#rm /opt/ngccli/*.zip
#chmod u+x /opt/ngccli/ngc## Running passed command
if [[ "$1" ]]; theneval "$@"
fi

2.0GA版本对应的是修改 install_ngc_cli.sh脚本

#!/usr/bin/env bash
set -eo pipefail# Select NGC CLI type based on command line arg
BATCH_CLI='ngccli_bat_linux.zip'
REG_CLI='ngccli_reg_linux.zip'# Installing NGC CLI type based on env variable.
if [ "x$NGC_INSTALL_CLI" == 'xBATCH' ]; thenCLI="$BATCH_CLI"
elif [ "x$NGC_INSTALL_CLI" == 'xREGISTRY' ]; thenCLI="$REG_CLI"
elseecho "Invalid NGC_INSTALL_CLI asked for. Exiting"exit 1
fi## Download and install
#mkdir -p /opt/ngccli && \
#wget "https://ngc.nvidia.com/downloads/$CLI" -P /opt/ngccli && \
#unzip -u "/opt/ngccli/$CLI" -d /opt/ngccli/ && \
#rm /opt/ngccli/*.zip && \
#chmod u+x /opt/ngccli/ngc## Running passed command
if [[ "$1" ]]; theneval "$@"
fi

2 准备数据集

接下来准备我们的训练数据，数据集的要求如下：

tlt-train不支持多分辨率图片的训练，需要将图片resize到统一大小，同时要对应的更改lable。官方建议图片的宽高应为16的整数倍，例如1920x1080

图片resize的方法可以参考下边的代码

def resize_img(img_path, save_path, img_size):w = img_size[0]h = img_size[1]img_list = os.listdir(img_path)for i in img_list:if i.endswith('.jpg'):img_array = cv2.imread((img_path + '/' + i), cv2.IMREAD_COLOR)new_array = cv2.resize(img_array, (w, h), interpolation=cv2.INTER_CUBIC)img_name = str(i)if os.path.exists(save_path):print(i)save_img = save_path + img_namecv2.imwrite(save_img, new_array)else:os.mkdir(save_path)save_img = save_path + img_namecv2.imwrite(save_img, new_array)

提供一个label resize的思路：先将xml中的xmin\ymin\xmax\ymax进行归一化保存，然后再乘以新的长宽即可

2.1 label转换

TLT的训练标注数据需要用KITTI数据集格式，所以如果你的数据集是其它格式的，需要做一个转化，这里我们提供VOC数据集和coco数据集的转化方法

2.1.1 voc数据集转换方法

import xml.etree.ElementTree as ET
import osbase_xml_dir = "./vocdata/Annotations/"
xml_list = os.listdir(base_xml_dir)
kitti_saved_dir = "./kitti/kitti_txt/"def convert_annotation(file_name):in_file = open(base_xml_dir + file_name)tree = ET.parse(in_file)root = tree.getroot()with open(kitti_saved_dir + file_name[:-4] + '.txt', 'w') as f:for obj in root.iter('object'):cls = obj.find('name').textxmlbox = obj.find('bndbox')xmin, ymin, xmax, ymax = xmlbox.find('xmin').text, xmlbox.find('ymin').text, \xmlbox.find('xmax').text, xmlbox.find('ymax').textf.write(cls + " " + '0' + " " + '0' + " " + '0' + " " + str(xmin) + '.0' + " "+ str(ymin) + '.0' + " " + str(xmax) + '.0' + " " + str(ymax) + '.0' + " " +'0' + " " + '0' + " " + '0' + " " + '0' + " " + '0' + " " + '0' + " " + '0' + '\n')for i in xml_list:convert_annotation(i)

2.1.2 coco数据集转换方法

def coco2kitti(coco_path, kitti_path, classes):width = 1280height = 720# 创建保存结果的文件夹if not os.path.exists(kitti_path):os.mkdir(kitti_path)for root, _, files in os.walk(coco_path):for file in files:filename, extension = os.path.splitext(file)print("------------------"+file+"------------------")if extension == '.txt':file_path = os.path.join(root, file)out_file = open(kitti_path + '%s.txt'%(filename), 'w')with open(file_path, "r") as fr:lines = fr.readlines()for line in lines:data = line.split(" ")try:class_id = int(data[0])x = float(data[1])y = float(data[2])w = float(data[3])h = float(data[4])except:print(file + 'is wrong')os.remove(file_path)# 删除有问题的label# os.remove(img_path + '%s.jpg'%(filename))os.remove(kitti_path + '%s.txt'%(filename))continueXmax = (2*x*width + w* width) / 2Xmin = (2*x*width - w* width) / 2Ymax = (2*y*height + h* height) / 2Ymin = (2*y*height - h* height) / 2bb = (round(Xmin, 2), round(Ymin, 2), round(Xmax, 2), round(Ymax, 2))print(bb)out_file.write(classes[class_id] + " 0 0 0 " + " ".join([str(a) for a in bb]) + " 0 0 0 0 0 0 0" + '\n')

2.2 创建目录

按照下边的目录结构创建各级目录（用下边的目录名字配置文件改动最少）

└── tlt-experiments
├── data
│   ├── faster_rcnn
│   ├── testing
│   │   └── image_2
│   └── training
│    ├── image_2
│    └── label_2
└── tfrecords
└── kitti_trainval

创建完文件夹后，将自己的训练图片和对应的KITTI标签放到training目录对应位置，在testing文件夹中放入测试图片

3 模型训练

TLT提供了很多jupyter notebook文件，将运行的命令都封装起来了，我们直接使用这些文件来进行训练就可以，当然也可以在容器中自己敲命令来训练。

进入容器中后，在workspace路径下，运行下边的启动命令

jupyter notebook --ip 0.0.0.0 --allow-root

然后根据token，在自己的浏览器中打开对应的IP和端口，打开后内容如下

3.1 环境设置

我们依次点击examples/faster_rcnn/faster_rcnn.ipynb，进入jupyter文件，如下

将第0步中KEY的值替换为自己的KEY值，获取方法参考1.1步骤，然后运行这个cell，就会输出运行结果

下边的第1步是下载KITTI数据集，因为我们准备了自己的数据集，这一步可以跳过不运行，为了防止出错，可以把cell中的代码全部注释掉，如下

3.2 生成tfrecords

接下来的步骤是生成tfrecords文件，首先需要修改tfrecords spec文件，这个文件里配置了数据集的位置信息

依次点击打开examples/faster_rcnn/specs/frcnn_tfrecords_kitti_trainval.txt文件，然后修改里边的值，如果你按照我上边的方法创建了文件夹，路径那些都不需要修改，如果你的图片是jpg格式，则只需要把imge_extension改一下，如下

其它都使用默认即可，修改完之后ctrl + s 保存一下，回到faster_rcnn的jupyter文件

依次运行接下来的两个cell

接下来你可以显示你生成的tfrecords文件

3.3 预训练模型

jupyter中提供了预训练模型的下载方法，但速度比较慢，我们注释掉，直接从NGC中手动下载

进入NGC官网，选择MODELS中的TLT Object Detection

进入后选择File Browser，选择自己要训练的预训练模型，这里我们选择的是RESNET 50，然后点击下载到本地

下载完后，拷贝到/workspace/tlt-experiments/data/faster_rcnn路径下，并重命名为resnet50.hdf5，至此预训练模型准备完毕

3.4 训练配置文件

打开faster_rcnn/specs下边的default_spec_resnet50.txt文件，从上到下依次有下边内容需要修改：

size_height_width中的height和width修改为自己图片的实际值
image_extension修改为jpg
target_class_mapping要修改为自己的类别，将key和value一一对应，比如我自己的类别是bicycle和motorbike，则改成下边的内容

target_class_mapping {
key: 'bicycle'
value: 'bicycle'
}
target_class_mapping {
key: 'motorbike'
value: 'motorbike'
}

output_image_width和output_image_height修改为自己图片的实际值
batch_size_per_gpu：根据自己显卡情况设置，如果显存不大，可以设成较小值，比如设成1或4
num_epochs：训练回合数，根据经验100-200可以有一个好的效果，太小了什么也检测不出来，我自己设成了200
freeze_bn：默认为True，当你的batch size大于16的时候可以设置为False，让BN层自动计算均值和偏差
checkpoint_interval：设置训练过程中保存模型的间隔，默认为1，可以根据自己需求调整
inference_config/model修改为/workspace/tlt-experiments/data/faster_rcnn/frcnn_kitti_resnet50.epoch200.tlt
evaluation_config/model修改为/workspace/tlt-experiments/data/faster_rcnn/frcnn_kitti_resnet50.epoch200.tlt

关于学习率learning rate：目前只支持soft_start_annealing_schedule这一种策略，参数如下

1. min_learning_rate: minimum learning late to be seen during the entire experiment.

2. max_learning_rate: maximum learning rate to be seen during the entire experiment

3. soft_start: Time to be lapsed before warm up ( expressed in percentage of progress between 0 and 1)

4. annealing: Time to start annealing the learning rate

参数含义如下（来源）：

Learning rate scheduler adjusts learning rate in the following 3 phases:
Phase 1: 0.0 <= progress < soft_start:
Starting from min_lr exponentially increase the learning rate to base_lr
Phase 2: soft_start <= progress < annealing_start:
Maintain the learning rate at base_lr
Phase 3: annealing_start <= progress <= 1.0:
Starting from base_lr exponentially decay the learning rate to min_lr

base_lr: Maximum learning rate
soft_start: The progress at which learning rate achieves base_lr when starting from min_lr
annealing_start: The progress at which learning rate starts to drop from base_lr to min_lr

归纳一下就是：先从最小学习率增加到最大学习率，然后保持最大学习率，最后从最大学习率减小到最小学习率。两个时间节点由soft_start和annealing这两个参数决定。

3.5 开始训练

使用tlt-train指令开始训练，注意要改成自己的配置文件，如果有多个GPU的话，可以指定使用的GPU个数，默认为1

为了避免出现显存报错问题，训练前先设置下环境变量：

$ export TF_FORCE_GPU_ALLOW_GROWTH=true

3.6 检测模型

经过漫长的等待，训练结束后，可以验证以下自己的模型效果

会输出每个类别的AP值以及mAP值，如果模型效果不错的话可以进行裁剪，如果mAP值很低，或者直接 mAP = 0的话，说明训练有问题，需要检查下自己的配置文件重新训练

4 模型优化

经过上述步骤，我们得到了一个检测效果不错的模型，但不会直接拿来部署，而是需要prune裁剪一下，然后第二次训练，得到最终模型

4.1 模型剪枝

使用tlt-prune指令对模型进行剪枝

其中，最主要的是调整 pth 阈值，这个阈值用来平衡模型的速度与精度，阈值越大，剪枝的越厉害，得到的模型越小，推理速度越快，但精度损失越严重，相反，阈值越小，裁剪完的模型越大，精度保持的越好，但速度提升没那么明显，可以根据自己的情况调整这个值，例如可以依次试试0.3、0.5、0.7等值。

4.2 重新训练

首先宝贝一份specs/default_spec_resnet18_retrain_spec.txt为default_spec_resnet50_retrain_spec.txt，然后打开做几处修改，和3.4中修改的内容类似，额外需要修改下边的内容

feature_extractor修改为resnet:50
num_epochs可以设置的比第一阶段小一些，我这里设置了100
对应的把inference_config和evaluation_config中的model修改为/workspace/tlt-experiments/data/faster_rcnn/frcnn_kitti_resnet50_retrain.epoch100.tlt
把全文resnet18的地方都替换为resnet50

然后可以开始训练

训练完之后，同样可以使用tlt-evaluate来计算一下模型的mAP，看看与未剪枝的模型相比是否有大的下降

4.3 可视化结果

可以使用tlt-infer工具对模型的检测效果进行可视化，运行后会输出两个文件夹，分别是faster_rcnn/inference_results_imgs_retrain和faster_rcnn/inference_dump_labels_retrain，用来保存检测结果图片和标签

当然也可以运行未剪枝的配置文件来检测未剪枝模型的检测效果

5 模型部署

TLT支持两种类型的模型导出，分别是etlt格式和engine格式，都可以用在deepstream中，但是engine是平台相关的，跟cuda、tensorrt等版本相关，所以如果部署的平台与训练的平台环境有差异，不建议直接导出engine使用，而是导出etlt模型，然后在deepstream中自动生成engine。

5.1 导出etlt模型

支持导出三种精度的模型，分别是FP32/FP16/INT8，如果要生成INT8模型，还需要生成一个量化表calibration cache file，具体方法如下

实际上，无论选择的是哪种类型（FP32/FP16/INT8），导出的模型都是相同大小的，类型都是FP32，只不过如果选择INT8的时候，会额外生成INT8 calibration table。（来源）

5.2 导出engine模型

也可以直接导出tensorrt的engine模型，方法如下

5.3 模型推理

可以使用上述的模型进行推理，计算mAP值，方法如下

5.4 deepstream部署

导出的这些模型可以直接在deepstream中进行部署，关于具体的部署方法，参考另一篇文章《DeepStream5.0系列之TLT模型调用》。

6 问题记录

6.1 XLA_GPU_JIT错误

tensorflow.python.framework.errors_impl.InvalidArgumentError: Invalid device ordinal value (5). Valid range is [0, 4].while setting up XLA_GPU_JIT device number 5

实际可用GPU与docker指定GPU不符，例如启动docker时指定--gpus all，而运行时第5张卡被占用无法使用，就会报错，

同样的，如果你指定一张不存在的GPU，比如只有gpu0,而你指定了--gpus 2则会报类似的错

6.2 类别问题

ValueError: Layer #206 (named "dense_class_td"), weight <tf.Variable 'dense_class_td/kernel:0' shape=(2048, 8) dtype=float32_ref> has shape (2048, 8), but the saved weight has shape (2048, 7).

6.3 OOM问题

tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[1,128,120,160] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc[[{{node b3_conv4_1_bn/batchnorm/mul_1}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.[[{{node loss/dense_regress_td_loss/cond_3/Min/Switch}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

这个是显存不足导致的，将batch设置的小一些即可，注意TLT的GPU默认从0号开始，所以指定多GPU的时候，显卡不要有其它任务，否则也容易报显存不足问题

6.4 datasize问题

tensorflow.python.framework.errors_impl.InvalidArgumentError: Invalid JPEG data or crop window, data size 688128

首先检查图片有没有问题，然后图片需要resize到统一大小，不支持不同分辨率的图片一起训练，同时要注意标签要一起修改

6.5 loss nan问题

训练开始时，如果报错Invalid loss, terminating training

Batch 3056: Invalid loss, terminating training
Batch 3056: Invalid loss, terminating training
3057/3762 [=======================>......] - ETA: 6:56 - loss: nan   Batch 3056: Invalid loss, terminating training

则需要调整学习率，有两种方法：

减小最大学习率 max_learning_rate
增大 soft_start，从最小学习率增大到最大学习率的速度变慢点，也可以起到减小学习率的作用

NVIDIA之TLT迁移学习训练自己的数据集相关推荐

基于NVIDIA GPUs的深度学习训练新优化
基于NVIDIA GPUs的深度学习训练新优化 New Optimizations To Accelerate Deep Learning Training on NVIDIA GPUs 不同行业采用 ...
pytorch1.7教程实验——迁移学习训练卷积神经网络进行图像分类
只是贴上跑通的代码以供参考学习参考网址:迁移学习训练卷积神经网络进行图像分类需要用到的数据集下载网址: https://download.pytorch.org/tutorial/hymenopt ...
NVIDIA RTX A6000深度学习训练基准
NVIDIA RTX A6000深度学习训练基准 2021年1月4日在本文中,我们对RTX A6000的PyTorch和TensorFlow培训性能进行了基准测试.我们将其与Tesla A100,V ...
Ubuntu18 安装nvidia驱动进行深度学习训练
Ubuntu18 安装nvidia驱动进行深度学习训练安装nvidia驱动驱动不兼容而造成系统循环登录的解决方案安装cuda 安装cudnn 下载anacoda3 找不到conda命令的解决方法 ...
迁移学习训练分类模型实践第一篇
迁移学习训练分类模型实践第一篇前言数据获取.预处理构建模型查看模型参数量和FLOPs 测试模型前言为了简洁,本文不包含任何训练过程,仅介绍处理数据.构建模型.使用随机初始化权重推断: 关于 ...
迁移学习训练集准确率一直上不去_可以提高你的图像识别模型准确率的7个技巧...
假定,你已经收集了一个数据集,建立了一个神经网络,并训练了您的模型. 但是,尽管你投入了数小时(有时是数天)的工作来创建这个模型,它还是能得到50-70%的准确率.这肯定不是你所期望的. 下面是一些提 ...
深度学习-Tensorflow2.2-预训练网络{7}-迁移学习基础针对小数据集-19
使用预训练网络(迁移学习) 预训练网络是一个保存好的之前已在大型数据集(大规模图像分类任务)上训练好的卷积神经网络如果这个原始数据集足够大且足够通用,那么预训练网络学到的特征的空间层次结构可以作为有 ...
全网都在讲迁移学习，可你会写代码了吗？收藏我这个，10分钟开始你的迁移学习训练
文章目录前言第一步加载预训练模型并修改类别数第二步选择模型所有层/最后一层进行反向传播优化探讨:如何确定模型最后一层的名字是什么方法一: 查询源代码方法二: 查询模型的子模块名字前言 ...
ResNet网络详解并使用pytorch搭建模型、并基于迁移学习训练
1.ResNet网络详解网络中的创新点: (1)超深的网络结构(突破1000层) (2)提出residual模块 (3)使用Batch Normalization加速训练(丢弃dropout) (1 ...

NVIDIA之TLT迁移学习训练自己的数据集