本文介紹如何實現一個快速的風格遷移專案，並一步一步地通過本地機器和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，這是一個可以用於圖片分類、背景分割的資料集。
在清楚了這些之後，我們即可正式開始我們的訓練。訓練過程可以分為以下幾步：
- 下載資料集檔案
- 下載網路模型檔案
- 編寫/下載訓練的程式碼
- 在本機訓練/使用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的訓練過程，主要可以分為以下幾步：
- 註冊賬號和實名驗證
- 整理檔案結構
- 上傳資料集到UFile
- 打包程式碼映象並上傳
- 建立訓練任務
- 檢視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中。