tensorflow 显存 训练_Nvidia 迁移学习训练(TLT)之旅

半个多月前听了英伟达何老师的在线TLT(transfer learning training 迁移学习)课程，一直想抽空实际操作一番，但由于工作忙碌，直到最近才有空去尝试TLT。现在是2020新年期间，但是肺炎盛行，全国各地严阵以待，不串门不聚餐，每天宅在家中，正好得空好好学习下此课程。

下图是TLT的流程，可以看出它主要有下面几个步骤。

• 基于预训练模型，进行TLT训练+部署。
• 使用增量数据，进行训练。
• 对上一步的模型进行裁剪优化，减少模型尺寸，最终能加速端侧推理。
• 对裁剪后的模型进行再训练。
• 输出到端侧进行推理。

但凡使用过Nvidia的端侧设备进行推理任务，例如Nano，TX2，Xavier的，都或多或少地遇到过自己的模型，在转换TensorRT engine的过程中，遇到有层不能转换成功的问题。英伟达对此的解决方法是书写Plugin层，但其规则多，而且我们对cuda编程不了解，往往不得不放弃。而使用TLT这个方法，我觉得会比较好的解决这个部署困难的问题，也就是使用迁移学习的方法，在预训练模型上，进行增量训练，得到的模型能直接部署到端侧，而且也能得到好的性能。

英伟达TLT提供的预训练模型有：

有Resnet，SSD，mobilenet 等主流模型，但没有包括一些最新的模型，期待英伟达能不断增加新模型的支持。

首先介绍下我的电脑环境

我使用的是台笔记本，里面有张1070的显卡，采用pascal架构，8G的GPU显存，安装了ubuntu 1804系统。能满足TLT的要求

TLT电脑硬件需求：

最低配置

4 GB system RAM

4 GB of GPU RAM

Single core CPU

1 GPU

50 GB of HDD space

推荐配置

32 GB system RAM

32 GB of GPU RAM

8 core CPU

4 GPUs

100 GB of SSD space

第二步TLT的环境预安装：

首先是安装Nvidia GPU的驱动。

添加nvidia 仓库

sudo add-apt-repository ppa:graphics-drivers/ppa

sudo apt update

显示可用的驱动版本

ubuntu-drivers devices

从上一步的版本列表中，选取最新的版本进行安装

sudo apt install nvidia-driver-440

其次就是安装docker-ce 社区版的docker，参照下面的页面指示安装。

https://docs.docker.com/install/linux/docker-ce/ubuntu/

然后安装nvidia-docker2

https://github.com/nvidia/nvidia-docker/wiki/Installation-(version-2.0)

最后就是到NGC 注册账户和获取API key，获取步骤如下图所示，猜测此key是用来解密模型用的密钥。

https://ngc.nvidia.com