UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

NO IMAGE

本文介紹如何實現一個快速的風格遷移項目,並一步一步地通過本地機器和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來供我們調用。

訓練

機器學習任務的訓練是一個從數據中提取特徵的過程,這一過程往往需要大量的數據集的支撐。在選擇數據集之前,不得不解釋一下風格遷移這個任務其實和圖片識別是有很大的相似性的。在最初的風格遷移中主要有兩個突破點,一是使用深度學習來獲得圖片的紋理,而之前的工作都是手工去建模的。另一方面,圖片中的紋理是通過局部的統計特徵來描述的。例如下圖:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

我們可以認為這張圖就是蘋果的紋理,用幾何來描述就可以說是[綠色或是紅色,有一個圓形的圖案,中間有一個點]這樣的描述。這些都可以認為是局部的特徵。而物體識別的任務中又大量的用到了局部的特徵識別器,如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 2014vgg_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即可。

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

如果想要實時可視化地看到訓練的過程,TensorFlow提供了非常好用的TensorBoard的組件。進入訓練過程中產生的models文件夾,使用tensorboard --logdir=mosaic,在瀏覽器中輸入127.0.0.1:6006即可查看TensorBoard中提供給我們的各項Summary,包括了各種的loss,如下圖:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

如果你使用的是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:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

在獲取到之後,我們還需要創建一個Docker的鏡像來存儲我們的代碼和數據。在上圖的菜單中選擇進入UHub,創建一個Docker鏡像即可。

第二步,我們來整理文件的結構,將我們的pretrained文件夾單獨移動出來並改名為data文件夾,並將原來的style-transfer文件夾命名為code文件夾(當然不改名也是可以的,但是目錄結構最好保持一致)。

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

第三步就需要將數據集上傳到UFile了。我們首先創建一個UFile的Bucket,由於我們上傳的文件較多(~58萬張),需要使用文件管理工具來上傳,因此需要在這裡選擇適合自己系統的文件管理工具下載。

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

下載之後,我們修改config.cfg的內容,將public_key和secret_key修改為我們自己的即可。之後我們即可運行

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

若結果如圖所示:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

則表明我們的數據集就開始上傳了,這可能會消耗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_pathtest_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)

之後,我們就可以運行我們的打包工具了。運行之前的腳本,若返回:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

就證明我們的打包過程是正確的。

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的內容複製進去即可。

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

第六步,查看訓練過程和結果。在訓練任務創建完成後,就可以在後臺看到訓練過程的日誌了。日誌會每5秒刷新一次,如果這個時候發現有什麼問題,可以停止任務,修改後重新運行。

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

在日誌中我們可以看到,UAI-Train會首先下載訓練所需的文件,之後啟動訓練過程。確保正常啟動後,就可以放下心來喝杯咖啡,休息一會兒再回來了:D

在任務執行期間,任務的右側會出現TensorBoard的按鈕,無需任何配置,點擊這個按鈕,我們就可以看到TensorBoard的詳情:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

在運行了大約40000步後,我們終於得到了我們想要的預訓練好的模型,此時的loss基本是收斂了。在完成之後我們的訓練程序會自動停止,並將模型保存在models文件夾中。在自己訓練時,可以使用TensorBoard來實時查看這一結果,如果loss不怎麼變化了,相對比較穩定了就可以停止了。

在訓練結束後,我們就可以看到訓練過程所花費的成本,並在UFile中看到對應的日誌和模型文件了:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

我稍微對比了在我的GTX 1070,單卡模式下的UAI-Train的時間和成本對比,以供參考。

環境每步平均用時成本總用時
GTX 10700.6單卡3000元左右4小時
UAI 單卡0.4730元左右3.2小時

執行

有了模型之後,再來執行就容易很多了。在代碼中找到eval.py文件,並使用如下的命令即可:

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

其中image_file可以指定為其他的圖片,就可以生成風格遷移後的圖片了。執行後我們可以看看這次輸出的所用時間:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

可以看到,基本上1秒左右就可以完成,相比於之前的每一張都需要接近10分鐘有了極大的提升。由於迭代步數的提高,效果也是非常棒:

UCloudAITrain深度學習實踐:使用TensorFlow實現快速風格遷移

如果對效果不滿意,還可以在conf/*.yml文件中修改對應的style_weight的值。如果紋理太多以至於找不到圖片的內容,就降低style_weight的值,反之,如果紋理太少,就可以提高style_weight的值。通過不斷地調整和嘗試來找到一個最佳的範圍。

需要注意的地方

在文件的輸入輸出的位置中,要注意從掛載點來讀寫。如模型的輸入路徑為/data/output,輸入路徑為/data/data/train2014。如果不符合的話,會導致無法讀取文件,或者無法將訓練好的模型上傳到UFile中。

相關文章

分佈式熔斷降級平臺aegis

開源Go,Py項目找隊友、顧問、貢獻者、贊助等等

如何在五分鐘內搭建人臉檢測/關鍵點識別等服務?

開源一個從標註到訓練、測試的計算機視覺庫