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的时间和成本对比,以供参考。

执行

有了模型之后，再来执行就容易很多了。在代码中找到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中。