裝逼一步到位！GauGAN代碼解讀來了
AI神筆馬良
如何裝逼一步到位？從涂鴉到栩栩如生，英偉達(dá)的神器GauGAN來了，擁有神筆馬良之手，五步畫馬就是這么簡單。下面讓我們一起解析GauGAN的代碼實踐過程。
GauGAN簡介
英偉達(dá)的“神筆馬良“，可以把涂鴉變成風(fēng)景畫，如同下圖：
這個模型可以將一張涂鴉（語義圖）生成相應(yīng)的一張真實圖，不少原畫師已經(jīng)用來畫背景圖了，它是怎么做到的？我們今天就來解讀GanGAN代碼，并使用相關(guān)數(shù)據(jù)集進(jìn)行實驗。
代碼解讀
開始前，先看下模型的創(chuàng)新思路：作者假設(shè)，如果語義圖輸入單個類別，傳統(tǒng)的歸一化層（normalization layer）傾向于“抹除”語義信息，因此作者提出了SPADE（SPatially-Adaptive (DE)normalization） 方法，GauGAN直接使用語義mask對輸出激活值進(jìn)行建模，即提出的空間自適應(yīng)層。并且其方法可以應(yīng)對各種使用語義圖的生成任務(wù)。
環(huán)境配置
將SPADE代碼克隆至本地，考慮到一些地區(qū)的網(wǎng)絡(luò)可能不太好，提供一個加速通道：
git clone https://github.com.cnpmjs.org/NVlabs/SPADE.git
進(jìn)入目錄，根據(jù)命令行安裝依賴和歸一化分支：
# 安裝依賴
cd SPADE/
pip install -r requirements.txt
# Synchronized BatchNorm
cd models/networks/
git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm .
cd ../../
數(shù)據(jù)集ADE20K準(zhǔn)備
下載地址：https://gas.graviti.cn/dataset/shannont/ADE20K?utm_medium=0728datawhale
數(shù)據(jù)規(guī)模：
訓(xùn)練集：25574；驗證集：2K；測試集：3K 內(nèi)容：365個場景；3688個類別。
ADE20K 的數(shù)據(jù)集，可用于場景感知、解析、分割、多物體識別和語義理解。
圖像涵蓋了SUN和Places數(shù)據(jù)集中的場景，官網(wǎng)可視化給出了目標(biāo)、部件和注釋示例的數(shù)量的列表。
樹狀表只展示了帶有超過 250 個注解示例的目標(biāo)，以及帶有超過 10 個注解示例的部件。
一些類別可以既是目標(biāo)，也是部件。例如，一個「門」可以是一個目標(biāo)（在一張室內(nèi)圖片中）或者一個部件（當(dāng)它是車的一個門時）。一些目標(biāo)經(jīng)常是部件（比如一條腿、一只手），盡管在某些情況下它們看起來與整體是相互獨立的（比如車庫中的汽車輪子）。
而有些目標(biāo)則永遠(yuǎn)不是部件（比如一個人，一輛卡車等等）。依據(jù)于部件所屬的目標(biāo)，相同的名稱類別（比如門）可對應(yīng)于若干個視覺范疇。例如，一個汽車的門從視覺上看是不同于一個櫥柜的門的。然而它們也共享一些相似的可供性（affordance）。
當(dāng)目標(biāo)不是另一個目標(biāo)的一個部件時，其分割掩碼將出現(xiàn)在 * _seg.png 內(nèi)。
如果分類是一個部件，則分割掩碼將出現(xiàn)在 * _seg_parts.png 內(nèi)。正確檢測目標(biāo)需要區(qū)分目標(biāo)是否表現(xiàn)為獨立目標(biāo)，或者是否是另一目標(biāo)的一個部件。
代碼結(jié)構(gòu)
官方提供的代碼結(jié)構(gòu)如下：
data：實現(xiàn)數(shù)據(jù)加載base_dataset.py是最底層的數(shù)據(jù)結(jié)構(gòu)，它繼承pytorch的data.Dataset類建立一個BaseDataset類，定義了get_params，get_transform等等一些函數(shù)，方便之后的使用。pix2pix_dataset.py創(chuàng)建了一個Pix2pixDataset類，繼承了BaseDataset類，這個類才是這個工程基本的數(shù)據(jù)類型，它重寫了getitem()函數(shù)，這個函數(shù)返回一個input_dict的字典，索引包括label，instance，image，path其他的.py文件都是自定義的數(shù)據(jù)集的類(ADE20K，coco等等)。
options：定義命令行參數(shù)，base_options.py定義公共的命令行參數(shù)，train_options.py和test_options.py分別定義訓(xùn)練和測試的命令行參數(shù)。base_options.py中包括一些在自定義數(shù)據(jù)集中比較重要的參數(shù)：label_nc為輸入標(biāo)簽數(shù)量，contain_dontcare_label是否包含不需要的標(biāo)簽。use_vae是否加入風(fēng)格控制，no_instance訓(xùn)練過程是否加入實例。model：該文件夾下的network定義了模型的各個部分normalization.py文件內(nèi)定義了SPADE歸一化層，為了方便實用，也創(chuàng)建了不使用SPADE的結(jié)構(gòu)，同時定義了VGG19結(jié)構(gòu)。base_network.py定義基本的網(wǎng)絡(luò)架構(gòu)，后面的所有網(wǎng)絡(luò)類型都是繼承它得到的，它的功能包括打印網(wǎng)絡(luò)的結(jié)構(gòu)，初始化各個參數(shù)。
訓(xùn)練新模型
如果是使用作者文檔中提到的數(shù)據(jù)集，可以使用下面這行代碼進(jìn)行訓(xùn)練：
python train.py --name [experiment_name] --dataset_mode facades --dataroot [path_to_facades_dataset]
python train.py --name [experiment_name] --dataset_mode coco --dataroot [path_to_coco_dataset]
其中experiment_name是自定義的名稱，dataroot后接數(shù)據(jù)集的路徑。
由于ade20k的數(shù)據(jù)集比較大，我也會使用小型一點的ade20k_outdoor進(jìn)行訓(xùn)練。將ade20k_outdoor數(shù)據(jù)集存儲在datasets文件夾下。
自定義數(shù)據(jù)集可以參考使用此段命令：
python train.py --name ade_outdoor --dataset_mode custom --label_dir ./datasets/ade20k_outdoor/train_label --image_dir ./datasets/ade20k_outdoor/train_img --label_nc 300 --contain_dontcare_label
完成后在checkpoints中將看到保存好的訓(xùn)練模型，終端打印出Training was successfully finished。
測試模型
接下來你可以使用剛剛訓(xùn)練好的模型進(jìn)行測試，name_of_experiment是上面自定義的名稱，
python test.py --name [name_of_experiment] --dataset_mode [dataset_mode] --dataroot [path_to_dataset]
自定義數(shù)據(jù)集使用下面這行：
python test.py --name [name_of_experiment] --dataset_mode custom --label_dir ./datasets/ade20k_outdoor/val_label --image_dir ./datasets/ade20k_outdoor/val_img --label_nc 300 --contain_dontcare_label
輸出結(jié)果將默認(rèn)保存至results文件夾。
結(jié)果分析
下面展示一下借助GauGAN由標(biāo)簽圖生成真實圖像的結(jié)果：
從生成結(jié)果來看，本文方法有更好的視覺效果，artifact也相對比較少?？梢哉J(rèn)為，SPADE的效果好主要還是更好的保留了語義圖中的語義信息的。