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中。