UCloud AI Train深度学习实践：使用TensorFlow实现快速风格迁移

本文介绍如何实现一个快速的风格迁移项目，并一步一步地经过本地机器和UCloud AI Train服务来实如今线的训练。在训练结束以后，对输入的图片进行风格迁移并输出成图片。本文也对比了使用本地机器和UCloud AI Train的时间/价格对比。

以前的一篇文章介绍了使用Pytorch来实现相似PRISMA的效果，然而评论区大佬们纷纷表示在本身的电脑上运行速度过慢，在当时个人机器上，一张图片动辄须要五六分钟的时间才能完成。这是因为每次咱们将两张图片输入到系统时，都要先对风格文件进行处理，训练出能够有效生成图片的网络，这一步是和目标图片无关的。若是咱们能够先生成这个模型，保存下来，以后每一次有新的图片进来，咱们就再也不从新训练模型了，这样就能够大大减小执行时间。

实际上，像PRISMA最初也是只给定了几种风格供选择，经过提早进行训练就能够大大节省服务器的成本，同时也减小了用户的等待时间。后续也有不少篇文章和代码来实现提早训练，如Perceptual Losses for Real-Time Style Transfer and Super-Resolution和知乎上的深度学习实践：使用Tensorflow实现快速风格迁移

此次，咱们就来按照这种方式来实现一个能够预训练的图片风格迁移的速度更快的应用。这部分的代码主体上采用的是知乎上@何之源大神的代码，只不过将网络从VGG 16更换到了VGG 19，而且用一些其余的网络进行了测试和对比。

对于多数机器学习问题，咱们均可以将其分为训练和执行两个部分。本文也分别对这两个部分进行解释。在训练部分传入风格的图片文件和指定的参数，使用TensorFlow生成模型的文件，在执行部分，则传入目标图片，并生成对应的转换后的图片。最后实现一个简单的API来供咱们调用。

训练

机器学习任务的训练是一个从数据中提取特征的过程，这一过程每每须要大量的数据集的支撑。在选择数据集以前，不得不解释一下风格迁移这个任务其实和图片识别是有很大的类似性的。在最初的风格迁移中主要有两个突破点，一是使用深度学习来得到图片的纹理，而以前的工做都是手工去建模的。另外一方面，图片中的纹理是经过局部的统计特征来描述的。例以下图：

咱们能够认为这张图就是苹果的纹理，用几何来描述就能够说是[绿色或是红色，有一个圆形的图案，中间有一个点]这样的描述。这些均可以认为是局部的特征。而物体识别的任务中又大量的用到了局部的特征识别器，如VGG 19。 因而咱们就能够考虑直接用物体识别的模型和数据集来进行局部纹理，或者说局部的特征的提取。更具体地，咱们使用了COCO Dataset，这是一个能够用于图片分类、背景分割的数据集。

在清楚了这些以后，咱们便可正式开始咱们的训练。训练过程能够分为如下几步：

1. 下载数据集文件
2. 下载网络模型文件
3. 编写/下载训练的代码
4. 在本机训练/使用UCloud AI Train服务来训练

首先咱们在这里(~13 GB)将数据集的压缩包下载下来。解压它会获得一个train 2014的文件夹。

除此以外，因为使用了VGG 19的网络，咱们也须要下载它的网络模型，能够在这里下载, 若是有兴趣实验一下其余的网络，也能够在Pretrained Models中寻找如ResNet、Inception等模型来尝试。本文使用的是VGG 19和VGG 16的网络。

我也已经把相关的数据集和模型文件放在了百度云，在国内下载速度应该会更快，能够在这里下载

我已经把训练的代码放在了github上，咱们首先将代码clone下来:

git clone https://github.com/stevefermi/style-transfer.git
复制代码

以后在该目录中新建一个名为pretrained的文件夹。将咱们下载获得的数据集train 2014和vgg_19.ckpt(或vgg_16.ckpt，相似)放入该文件夹。 最终获得的文件夹结构以下：

- style transfer
 - pretrained
     - train2014
     - vgg_19.ckpt
 - train.py
 - eval.py
 - conf
 - others...
复制代码

以后在style transfer这一级目录中执行python train.py便可开始训练了。默认会选择conf/mosaic.yml做为风格样式的图片。若是想训练不一样的风格样式的图片，经过-c来指定：python train.py -c conf/candy.yml便可。

若是想要实时可视化地看到训练的过程，TensorFlow提供了很是好用的TensorBoard的组件。进入训练过程当中产生的models文件夹，使用tensorboard --logdir=mosaic，在浏览器中输入127.0.0.1:6006便可查看TensorBoard中提供给咱们的各项Summary，包括了各类的loss，以下图：

若是你使用的是NVIDIA的较强的显卡（大约GTX 1060及以上便可），就能够开启使用CUDA加速，这能够大大减小计算时间。以个人GTX 1070+16G内存+i7 6700的处理器，每步大约须要0.6秒左右。想要获得一个比较好的结果大约须要10~12个小时。若是实在等不及的话，能够直接下载训练好的模型文件。固然，随着如今云计算的发展，在服务器上进行训练成为了一个成本更低，性价比更高的方案。咱们接下来就尝试使用UCloud提供的UAI-Train服务来进行，UAI提供了P40的单卡/四卡方案，价格每分钟须要0.17/0.6元，相比于本身购置机器来跑训练，成本下降了许多。

在UCloud上进行UAI的训练过程，主要能够分为如下几步：

1. 注册帐号和实名验证
2. 整理文件结构
3. 上传数据集到UFile
4. 打包代码镜像并上传
5. 建立训练任务
6. 查看TensorBoard，下载训练好的模型文件

第一步，注册帐号和实名验证。首先注册一个UCloud的帐号，在注册和实名验证以后，便可在以下图所示位置获取Public和Private Key：

在获取到以后，咱们还须要建立一个Docker的镜像来存储咱们的代码和数据。在上图的菜单中选择进入UHub，建立一个Docker镜像便可。

第二步，咱们来整理文件的结构，将咱们的pretrained文件夹单独移动出来并更名为data文件夹，并将原来的style-transfer文件夹命名为code文件夹（固然不更名也是能够的，可是目录结构最好保持一致）。

第三步就须要将数据集上传到UFile了。咱们首先建立一个UFile的Bucket，因为咱们上传的文件较多（~58万张），须要使用文件管理工具来上传，所以须要在这里选择适合本身系统的文件管理工具下载。

下载以后，咱们修改config.cfg的内容，将public_key和secret_key修改成咱们本身的便可。以后咱们便可运行

./filemgr-win64.exe --action mput --bucket demobucket --key key --dir train2014 --prefix train2014/ --trimpath <path to train2014 folder>  
复制代码

若结果如图所示：

则代表咱们的数据集就开始上传了，这可能会消耗30分钟到一个小时，取决于网速。因为咱们将VGG模型的checkpoints(ckpt)文件打包在了代码的Docker镜像中，所以这一步就再也不须要上传这一模型文件了。

第四步咱们来打包代码镜像并将它上传。这一步是相当重要的一步。

4.1 打包上传须要UAI-Train服务的SDK，在确保已经安装了git以后，能够在一个目录中使用以下命令：

git clone https://github.com/ucloud/uai-sdk.git
cd uai-sdk
sudo python setup.py install # 视状况决定需不须要使用sudo
复制代码

并将uaitrain_tool/tf/tf_tool.py的文件复制到和code同级的文件夹中，最终的代码和数据的结构以下：

- code
    - train.py
    - others...
- data
    - train2014
    - others...
- tf_tool.py
复制代码

4.2 开始进行打包。在放置好tf_tool.py并整理好文件结构后，咱们就能够开始进行打包了。咱们能够建立一个.bat(windows)文件或者.sh(Linux)脚原本方便咱们打包，内容以下：

python tf_tool.py pack --public_key=<public_key> --private_key=<private_key> --code_path=./code/ --mainfile_path=train.py --uhub_username="<username>" --uhub_password="<password>" --uhub_registry=<your_image_name> --uhub_imagename=<your_image_name> --uhub_imagetag=<your_image_tag> --internal_uhub=false --ai_arch_v=tensorflow-1.3.0 --test_data_path=./data --test_output_path=./data/models --train_params="" --python_version=python-3.6.2 --os=ubuntu-16.04 
复制代码

这其中test_data_path和test_output_path是指在本地测试时的数据集和导出的模型文件的位置，建议使用绝对路径。其中须要替换public/private key为本身以前获取到的私钥，将uhub-registry更换为本身的镜像名称。uhub_username和password就是UCloud的帐号密码。image_tag是标签文件，能够写入版本号等信息

UCloud的打包工具会将咱们传入的data这一文件夹挂载在容器的/data/data文件夹下，所以也要对应的修改咱们的训练和配置文件中的路径。分别是train.py中第39行，将pretrained/train2014修改成/data/data/train2014，以及conf/mosaic.yml中的loss_model_file: "pretrained/vgg_19.ckpt" # the path to the checkpoint修改成loss_model_file: "/data/data/vgg_19.ckpt" # the path to the checkpoint

而且，值得特别注意的是，记得要修改模型保存的路径为/data/output。在咱们的代码中，也就是第23行的training_path改成training_path = os.path.join('/data/output', FLAGS.naming)

以后，咱们就能够运行咱们的打包工具了。运行以前的脚本，若返回：

就证实咱们的打包过程是正确的。

4.3 查看打包结果。在打包结束以后，咱们能够看到在目录下自动生成了uaitrain_cmd.txt等文件。在这个文件中包含了本地测试和部署的命令，在运行结束以后，生成了下面的文件：

CMD Used for deploying: /data/train.py 
CMD for CPU local test: sudo docker run -it -v D:\Learn\tensorflow\uai-prisma\data:/data/data -v D:\Learn\tensorflow\uai-prisma\data\models:/data/output style_transfer-cpu:v0.1 /bin/bash -c "cd /data && /usr/bin/python /data/train.py --work_dir=/data --data_dir=/data/data --output_dir=/data/output --log_dir=/data/output"
CMD for GPU local test: sudo nvidia-docker run -it -v D:\Learn\tensorflow\uai-prisma\data:/data/data -v D:\Learn\tensorflow\uai-prisma\data\models:/data/output uhub.ucloud.cn/style_transfer/style_transfer:v0.1 /bin/bash -c "cd /data && /usr/bin/python /data/train.py --num_gpus=1 --work_dir=/data --data_dir=/data/data --output_dir=/data/output --log_dir=/data/output"

复制代码

第五步，建立AI训练任务。回到UCloud的控制台，选择AI训练服务。在建立AI训练任务时，因为须要在Ufile中下载数据，因此须要填入地址和Public/Secret Key，并将uaitrain_cmd.txt这一文件中的CMD Used for deploying的内容复制进去便可。

第六步，查看训练过程和结果。在训练任务建立完成后，就能够在后台看到训练过程的日志了。日志会每5秒刷新一次，若是这个时候发现有什么问题，能够中止任务，修改后从新运行。

在日志中咱们能够看到，UAI-Train会首先下载训练所需的文件，以后启动训练过程。确保正常启动后，就能够放下心来喝杯咖啡，休息一下子再回来了:D

在任务执行期间，任务的右侧会出现TensorBoard的按钮，无需任何配置，点击这个按钮，咱们就能够看到TensorBoard的详情：

在运行了大约40000步后，咱们终于获得了咱们想要的预训练好的模型，此时的loss基本是收敛了。在完成以后咱们的训练程序会自动中止，并将模型保存在models文件夹中。在本身训练时，可使用TensorBoard来实时查看这一结果，若是loss不怎么变化了，相对比较稳定了就能够中止了。

在训练结束后，咱们就能够看到训练过程所花费的成本，并在UFile中看到对应的日志和模型文件了：

我稍微对比了在个人GTX 1070，单卡模式下的UAI-Train的时间和成本对比,以供参考。

环境	每步平均用时	成本	总用时
GTX 1070	0.6	单卡3000元左右	4小时
UAI 单卡	0.47	30元左右	3.2小时

执行

有了模型以后，再来执行就容易不少了。在代码中找到eval.py文件，并使用以下的命令便可：

python eval.py --model_file .\models\mosaic\fast-style-model.ckpt-done --image_file .\img\test.jpg
复制代码

其中image_file能够指定为其余的图片，就能够生成风格迁移后的图片了。执行后咱们能够看看此次输出的所用时间：

能够看到，基本上1秒左右就能够完成，相比于以前的每一张都须要接近10分钟有了极大的提高。因为迭代步数的提升，效果也是很是棒：

若是对效果不满意，还能够在conf/*.yml文件中修改对应的style_weight的值。若是纹理太多以致于找不到图片的内容，就下降style_weight的值，反之，若是纹理太少，就能够提升style_weight的值。经过不断地调整和尝试来找到一个最佳的范围。

须要注意的地方

在文件的输入输出的位置中，要注意从挂载点来读写。如模型的输入路径为/data/output，输入路径为/data/data/train2014。若是不符合的话，会致使没法读取文件，或者没法将训练好的模型上传到UFile中。