戰(zhàn)：從論文到代碼的神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索避坑指南)
1. 為什么值得花時(shí)間啃透DARTS如果你正在做神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索NAS方向的研究或者工程落地大概率繞不開(kāi)DARTS這篇工作。全稱Differentiable Architecture Search核心思路一句話就能說(shuō)清楚把離散的架構(gòu)選擇松弛成連續(xù)的、可微的權(quán)重然后用梯度下降同時(shí)優(yōu)化網(wǎng)絡(luò)權(quán)重和架構(gòu)參數(shù)。這個(gè)想法在2018年提出來(lái)的時(shí)候直接炸了NAS這個(gè)圈子因?yàn)樗言拘枰獛浊PU天的搜索過(guò)程壓縮到了幾個(gè)GPU天。我自己第一次讀這篇論文的時(shí)候感覺(jué)公式推導(dǎo)很順代碼應(yīng)該也不難。結(jié)果真正把官方代碼拉下來(lái)跑通、再自己改結(jié)構(gòu)的時(shí)候才發(fā)現(xiàn)坑遠(yuǎn)比想象中多。從搜索階段的顯存爆炸到離散化之后的性能崩塌再到評(píng)估階段的復(fù)現(xiàn)差異每一步都有大量論文里不會(huì)寫(xiě)的細(xì)節(jié)。這篇文章就是把我從讀論文到跑代碼、再到自己改結(jié)構(gòu)這一整條鏈路上積累的東西全部攤開(kāi)講包括核心原理的拆解、代碼關(guān)鍵模塊的逐行分析、實(shí)操中遇到的典型問(wèn)題和排查思路。適合的讀者有三類一是剛接觸NAS、想通過(guò)DARTS入門(mén)的研究生或者算法工程師二是已經(jīng)跑過(guò)官方代碼但結(jié)果對(duì)不上、想搞清楚哪里出了問(wèn)題的同學(xué)三是想把DARTS的思想遷移到自己業(yè)務(wù)場(chǎng)景里的從業(yè)者。不管你是哪一類我都會(huì)盡量把“為什么這么做”講透而不是只貼一段代碼讓你抄。2. DARTS的核心設(shè)計(jì)思路拆解2.1 從離散搜索到連續(xù)松弛這個(gè)彎是怎么拐過(guò)來(lái)的傳統(tǒng)的NAS方法比如基于強(qiáng)化學(xué)習(xí)或者進(jìn)化算法的方案本質(zhì)上是在一個(gè)離散的搜索空間里做黑盒優(yōu)化。搜索空間里每個(gè)候選架構(gòu)都是一個(gè)獨(dú)立的點(diǎn)你需要完整訓(xùn)練它、評(píng)估它才能知道它好不好。這就導(dǎo)致搜索成本極高因?yàn)槊看卧u(píng)估都意味著一次完整的訓(xùn)練過(guò)程。DARTS的關(guān)鍵洞察在于我不需要把每個(gè)候選架構(gòu)當(dāng)成獨(dú)立的點(diǎn)我可以把搜索空間表示成一張有向無(wú)環(huán)圖每條邊上的操作是一個(gè)混合操作混合的權(quán)重就是架構(gòu)參數(shù)。這樣一來(lái)架構(gòu)搜索就變成了對(duì)這些連續(xù)權(quán)重的優(yōu)化問(wèn)題可以直接用梯度下降來(lái)解。具體來(lái)說(shuō)假設(shè)一個(gè)單元cell里有若干節(jié)點(diǎn)節(jié)點(diǎn)之間的邊上有若干候選操作比如3x3卷積、5x5卷積、跳連、池化等。DARTS的做法是在每條邊上放一個(gè)softmax加權(quán)的混合操作混合操作 sum over o of (exp(alpha_o) / sum over o of exp(alpha_o)) * o(x)其中alpha就是架構(gòu)參數(shù)。搜索的過(guò)程就是同時(shí)優(yōu)化網(wǎng)絡(luò)權(quán)重w和架構(gòu)參數(shù)alpha。訓(xùn)練完成后每條邊保留softmax權(quán)重最大的那個(gè)操作得到一個(gè)離散的架構(gòu)。這個(gè)松弛的好處是顯而易見(jiàn)的搜索過(guò)程變成了可微的可以用標(biāo)準(zhǔn)的反向傳播來(lái)優(yōu)化效率提升了好幾個(gè)數(shù)量級(jí)。但代價(jià)也很明顯松弛之后的連續(xù)空間和離散空間之間存在gap搜索時(shí)最優(yōu)的連續(xù)架構(gòu)離散化之后不一定還是最優(yōu)的。這個(gè)問(wèn)題在后面會(huì)詳細(xì)展開(kāi)。2.2 雙層優(yōu)化為什么不能簡(jiǎn)單地把a(bǔ)lpha和w一起訓(xùn)DARTS的優(yōu)化目標(biāo)是一個(gè)雙層優(yōu)化問(wèn)題bilevel optimization。外層優(yōu)化架構(gòu)參數(shù)alpha目標(biāo)是讓在驗(yàn)證集上的損失最小內(nèi)層優(yōu)化網(wǎng)絡(luò)權(quán)重w目標(biāo)是讓在訓(xùn)練集上的損失最小。用公式表示就是min over alpha of L_val(w*(alpha), alpha) s.t. w*(alpha) argmin over w of L_train(w, alpha)為什么不能直接把a(bǔ)lpha和w放在一起用訓(xùn)練集損失來(lái)優(yōu)化因?yàn)槟菢訒?huì)導(dǎo)致架構(gòu)參數(shù)過(guò)擬合訓(xùn)練集搜索出來(lái)的架構(gòu)在驗(yàn)證集上表現(xiàn)很差。雙層優(yōu)化的設(shè)計(jì)就是為了讓架構(gòu)參數(shù)在驗(yàn)證集上做選擇保證泛化性。但雙層優(yōu)化有個(gè)麻煩w*(alpha)本身是內(nèi)層優(yōu)化的結(jié)果要對(duì)alpha求導(dǎo)需要用到隱函數(shù)定理或者近似。DARTS采用的是一種一階近似的方法通過(guò)一步梯度下降來(lái)近似w*(alpha)然后對(duì)alpha求導(dǎo)。這個(gè)近似在論文里有詳細(xì)的推導(dǎo)代碼里也有對(duì)應(yīng)的實(shí)現(xiàn)。實(shí)際操作中這個(gè)近似會(huì)帶來(lái)一些問(wèn)題。比如搜索階段的batch size不能太小否則梯度的方差會(huì)很大導(dǎo)致alpha的優(yōu)化不穩(wěn)定。另外訓(xùn)練集和驗(yàn)證集的劃分也很關(guān)鍵如果驗(yàn)證集太小或者分布和訓(xùn)練集差異太大搜索出來(lái)的架構(gòu)可能會(huì)偏向于某些特定的操作。2.3 搜索空間設(shè)計(jì)cell結(jié)構(gòu)的選擇與權(quán)衡DARTS的搜索空間是基于cell的。一個(gè)cell是一個(gè)有向無(wú)環(huán)圖包含若干個(gè)節(jié)點(diǎn)和邊。DARTS定義了兩種cellnormal cell和reduction cell。Normal cell保持特征圖的空間分辨率不變r(jià)eduction cell將空間分辨率減半。每個(gè)cell有兩個(gè)輸入節(jié)點(diǎn)分別是前兩層的輸出和一個(gè)輸出節(jié)點(diǎn)。中間的節(jié)點(diǎn)通過(guò)邊連接每條邊上的操作從預(yù)定義的候選操作集合中選擇。候選操作通常包括3x3深度可分離卷積、5x5深度可分離卷積、3x3空洞卷積、5x5空洞卷積、3x3最大池化、3x3平均池化、identity跳連、zero無(wú)連接。這個(gè)設(shè)計(jì)的好處是搜索空間相對(duì)緊湊搜索出來(lái)的cell可以堆疊成任意深度的網(wǎng)絡(luò)。但也有一些限制比如cell的結(jié)構(gòu)是固定的只有邊上的操作是可搜索的節(jié)點(diǎn)的連接方式也是預(yù)定義的不能搜索。這些限制在一定程度上約束了搜索空間的表達(dá)能力但換來(lái)了搜索效率的大幅提升。我在實(shí)際使用中發(fā)現(xiàn)候選操作集合的選擇對(duì)搜索結(jié)果影響很大。如果池化操作太多搜索出來(lái)的架構(gòu)會(huì)傾向于大量使用池化導(dǎo)致模型容量不足如果卷積操作太多搜索又會(huì)偏向于堆疊卷積參數(shù)量爆炸。通常的做法是保持一個(gè)相對(duì)平衡的候選集合并根據(jù)具體任務(wù)做微調(diào)。3. 代碼關(guān)鍵模塊逐行解析3.1 搜索空間的構(gòu)建從代碼角度看cell是怎么定義的官方代碼里搜索空間的構(gòu)建主要在model_search.py這個(gè)文件里。核心類是Cell它接收一個(gè)steps參數(shù)表示cell里有多少個(gè)中間節(jié)點(diǎn)和一個(gè)multiplier參數(shù)表示輸出節(jié)點(diǎn)由多少個(gè)中間節(jié)點(diǎn)拼接而成。每個(gè)cell的構(gòu)建過(guò)程大致是這樣的首先定義兩個(gè)輸入節(jié)點(diǎn)前兩層的輸出然后對(duì)于每個(gè)中間節(jié)點(diǎn)從前面的所有節(jié)點(diǎn)中選擇兩個(gè)作為輸入在這兩條邊上分別放置混合操作?；旌喜僮鞯臋?quán)重就是架構(gòu)參數(shù)alpha。代碼里有一個(gè)關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)self._ops它是一個(gè)ModuleList每個(gè)元素是一個(gè)MixedOp對(duì)應(yīng)一條邊上的混合操作。MixedOp內(nèi)部維護(hù)了一個(gè)ModuleList包含所有候選操作以及一個(gè)alpha參數(shù)用于計(jì)算softmax權(quán)重。class MixedOp(nn.Module): def __init__(self, C, stride): super(MixedOp, self).__init__() self._ops nn.ModuleList() for primitive in PRIMITIVES: op OPS[primitive](C, stride, False) self._ops.append(op) def forward(self, x, weights): return sum(w * op(x) for w, op in zip(weights, self._ops))這段代碼看起來(lái)很簡(jiǎn)單但有幾個(gè)細(xì)節(jié)值得注意。第一weights是通過(guò)對(duì)alpha做softmax得到的這個(gè)softmax是在forward里計(jì)算的而不是在__init__里。第二sum這里用的是Python內(nèi)置的sum而不是torch.sum這意味著每個(gè)操作的輸出是逐個(gè)相加的顯存占用會(huì)比較高。第三OPS是一個(gè)字典把操作名映射到具體的nn.Module這個(gè)字典的定義在文件開(kāi)頭。在實(shí)際使用中我發(fā)現(xiàn)MixedOp的顯存占用是一個(gè)大問(wèn)題。因?yàn)槊織l邊上都要計(jì)算所有候選操作的輸出然后加權(quán)求和所以顯存占用是候選操作數(shù)量的倍數(shù)。如果候選操作有8個(gè)那顯存占用就是普通網(wǎng)絡(luò)的8倍左右。這也是為什么DARTS搜索階段需要比較大的顯存。3.2 架構(gòu)參數(shù)的初始化與更新alpha的優(yōu)化細(xì)節(jié)架構(gòu)參數(shù)alpha的初始化方式對(duì)搜索結(jié)果有一定影響。官方代碼里alpha是用torch.nn.Parameter定義的初始化為0。這意味著初始時(shí)所有操作的softmax權(quán)重是均勻的每條邊上的混合操作是所有候選操作的等權(quán)平均。alpha的更新是通過(guò)優(yōu)化器來(lái)完成的。官方代碼里用的是torch.optim.Adam學(xué)習(xí)率設(shè)置為3e-4權(quán)重衰減設(shè)置為1e-3。這個(gè)學(xué)習(xí)率比網(wǎng)絡(luò)權(quán)重的學(xué)習(xí)率通常0.025小很多原因是架構(gòu)參數(shù)需要更精細(xì)的調(diào)整學(xué)習(xí)率太大會(huì)導(dǎo)致搜索過(guò)程震蕩。在train_search.py里有一個(gè)關(guān)鍵的步驟是計(jì)算架構(gòu)參數(shù)的梯度。由于雙層優(yōu)化的存在這個(gè)梯度不能直接用驗(yàn)證集損失對(duì)alpha求導(dǎo)得到而是需要先對(duì)網(wǎng)絡(luò)權(quán)重做一步虛擬更新然后再計(jì)算梯度。代碼里對(duì)應(yīng)的邏輯是# 計(jì)算虛擬更新后的網(wǎng)絡(luò)權(quán)重 unrolled_model self._compute_unrolled_model(input, target, eta, network_optimizer) # 計(jì)算架構(gòu)參數(shù)的梯度 unrolled_loss self._loss(unrolled_model, self._steps, input_valid, target_valid) unrolled_loss.backward()這里的_compute_unrolled_model就是做一步虛擬更新eta是虛擬更新的學(xué)習(xí)率。這個(gè)步驟是DARTS的核心也是計(jì)算量最大的部分因?yàn)樗枰Ａ粲?jì)算圖顯存占用會(huì)進(jìn)一步增加。我在實(shí)際使用中發(fā)現(xiàn)如果顯存不夠可以把unrolled設(shè)置為False改用一階近似。這樣雖然搜索精度會(huì)有所下降但顯存占用會(huì)大幅降低。對(duì)于顯存有限的場(chǎng)景這是一個(gè)很實(shí)用的折中方案。3.3 離散化與架構(gòu)導(dǎo)出從連續(xù)權(quán)重到最終網(wǎng)絡(luò)搜索完成后需要把連續(xù)權(quán)重離散化得到最終的架構(gòu)。這個(gè)過(guò)程在model_search.py的genotype方法里實(shí)現(xiàn)。核心邏輯是對(duì)于每條邊選擇softmax權(quán)重最大的操作對(duì)于每個(gè)中間節(jié)點(diǎn)選擇權(quán)重最大的兩條邊作為輸入。def genotype(self): def _parse(weights): gene [] n 2 start 0 for i in range(self._steps): end start n W weights[start:end].copy() edges sorted(range(i 2), keylambda x: -max(W[x][k] for k in range(len(W[x])) if k ! PRIMITIVES.index(none)))[:2] for j in edges: k_best None for k in range(len(W[j])): if k ! PRIMITIVES.index(none): if k_best is None or W[j][k] W[j][k_best]: k_best k gene.append((PRIMITIVES[k_best], j)) start end n 1 return gene # ...這段代碼有幾個(gè)細(xì)節(jié)需要注意。第一none操作被排除在外也就是說(shuō)如果某條邊的none操作權(quán)重最大這條邊會(huì)被忽略。第二每個(gè)中間節(jié)點(diǎn)選擇兩條輸入邊這是DARTS的默認(rèn)設(shè)置。第三gene是一個(gè)列表每個(gè)元素是一個(gè)(操作名, 輸入節(jié)點(diǎn)索引)的元組。離散化之后通常會(huì)重新訓(xùn)練一個(gè)網(wǎng)絡(luò)來(lái)評(píng)估架構(gòu)的性能。這個(gè)重新訓(xùn)練的過(guò)程和普通網(wǎng)絡(luò)訓(xùn)練類似但有一些細(xì)節(jié)需要注意比如學(xué)習(xí)率調(diào)度、數(shù)據(jù)增強(qiáng)、正則化等。官方代碼里用的是標(biāo)準(zhǔn)的訓(xùn)練流程但在實(shí)際使用中這些超參數(shù)對(duì)最終性能影響很大。4. 實(shí)操過(guò)程中的關(guān)鍵環(huán)節(jié)與避坑指南4.1 環(huán)境配置與數(shù)據(jù)準(zhǔn)備從零開(kāi)始跑通官方代碼先把環(huán)境搭起來(lái)。官方代碼基于PyTorch建議用1.0以上的版本但不要用太新的版本因?yàn)橛行〢PI可能不兼容。Python版本建議3.6到3.8之間。CUDA版本根據(jù)你的顯卡來(lái)選但要注意PyTorch版本和CUDA版本的對(duì)應(yīng)關(guān)系。數(shù)據(jù)準(zhǔn)備方面官方代碼支持CIFAR-10和ImageNet。CIFAR-10可以直接用torchvision.datasets.CIFAR10下載ImageNet需要自己準(zhǔn)備。數(shù)據(jù)增強(qiáng)方面官方代碼用了標(biāo)準(zhǔn)的隨機(jī)裁剪和隨機(jī)翻轉(zhuǎn)另外還用了cutout。Cutout的實(shí)現(xiàn)很簡(jiǎn)單就是在圖像上隨機(jī)挖一個(gè)16x16的方塊把它置零。這個(gè)操作對(duì)防止過(guò)擬合很有效。class Cutout(object): def __init__(self, length): self.length length def __call__(self, img): h, w img.size(1), img.size(2) mask np.ones((h, w), np.float32) y np.random.randint(h) x np.random.randint(w) y1 np.clip(y - self.length // 2, 0, h) y2 np.clip(y self.length // 2, 0, h) x1 np.clip(x - self.length // 2, 0, w) x2 np.clip(x self.length // 2, 0, w) mask[y1: y2, x1: x2] 0. img torch.from_numpy(img.numpy() * mask) return img這段代碼看起來(lái)簡(jiǎn)單但有一個(gè)坑img.numpy()要求img是CPU上的tensor如果數(shù)據(jù)加載器返回的是CUDA tensor這里會(huì)報(bào)錯(cuò)。所以cutout要在數(shù)據(jù)加載階段做而不是在訓(xùn)練循環(huán)里做。4.2 搜索階段的參數(shù)設(shè)置與顯存優(yōu)化搜索階段的參數(shù)設(shè)置很關(guān)鍵。官方代碼里搜索階段訓(xùn)練50個(gè)epochbatch size設(shè)置為64網(wǎng)絡(luò)權(quán)重的學(xué)習(xí)率是0.025架構(gòu)參數(shù)的學(xué)習(xí)率是3e-4。這些參數(shù)在CIFAR-10上表現(xiàn)不錯(cuò)但換到其他數(shù)據(jù)集可能需要調(diào)整。顯存優(yōu)化是搜索階段的一個(gè)大問(wèn)題。因?yàn)槊織l邊上都要計(jì)算所有候選操作的輸出顯存占用是普通網(wǎng)絡(luò)的數(shù)倍。如果顯存不夠可以嘗試以下幾個(gè)方案減小batch size。這是最直接的方法但會(huì)影響搜索穩(wěn)定性。減少候選操作數(shù)量。比如去掉一些不常用的操作只保留核心的幾個(gè)。使用一階近似。把unrolled設(shè)置為False可以大幅降低顯存占用。使用混合精度訓(xùn)練。PyTorch的amp模塊可以自動(dòng)做混合精度顯存占用可以減少30%到50%。我自己的經(jīng)驗(yàn)是如果顯存只有8GB建議用一階近似加混合精度batch size設(shè)置為32左右。如果顯存有16GB以上可以用二階近似batch size設(shè)置為64。4.3 評(píng)估階段的重新訓(xùn)練與性能復(fù)現(xiàn)搜索完成后會(huì)得到一個(gè)離散的架構(gòu)。這個(gè)架構(gòu)需要重新訓(xùn)練來(lái)評(píng)估性能。官方代碼里評(píng)估階段訓(xùn)練600個(gè)epochbatch size設(shè)置為96學(xué)習(xí)率用余弦退火調(diào)度初始學(xué)習(xí)率0.025。評(píng)估階段有幾個(gè)坑需要注意。第一數(shù)據(jù)增強(qiáng)的強(qiáng)度要夠否則容易過(guò)擬合。除了標(biāo)準(zhǔn)的隨機(jī)裁剪和翻轉(zhuǎn)還可以加cutout、mixup、cutmix等。第二學(xué)習(xí)率調(diào)度很關(guān)鍵余弦退火比階梯下降效果更好。第三權(quán)重衰減要設(shè)置合理太大或太小都會(huì)影響性能。我在復(fù)現(xiàn)官方結(jié)果的時(shí)候發(fā)現(xiàn)評(píng)估階段的性能波動(dòng)很大。同樣的架構(gòu)不同的隨機(jī)種子最終精度可能差0.5%到1%。所以如果要對(duì)比不同架構(gòu)的性能最好用多個(gè)隨機(jī)種子跑取平均值。另外官方代碼里的評(píng)估階段用的是CIFAR-10的完整訓(xùn)練集但搜索階段用的是訓(xùn)練集的50%作為訓(xùn)練集另外50%作為驗(yàn)證集。這個(gè)劃分方式在評(píng)估階段要改回來(lái)否則會(huì)浪費(fèi)一半的訓(xùn)練數(shù)據(jù)。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 搜索階段loss不下降或者震蕩怎么辦搜索階段loss不下降或者震蕩通常有幾個(gè)原因。第一學(xué)習(xí)率設(shè)置不當(dāng)。架構(gòu)參數(shù)的學(xué)習(xí)率如果太大會(huì)導(dǎo)致搜索過(guò)程震蕩如果太小搜索會(huì)停滯。建議從3e-4開(kāi)始試如果震蕩就降到1e-4如果停滯就升到1e-3。第二數(shù)據(jù)劃分不合理。如果驗(yàn)證集太小或者分布和訓(xùn)練集差異太大架構(gòu)參數(shù)的梯度會(huì)很不穩(wěn)定。建議驗(yàn)證集至少占訓(xùn)練集的30%到50%。第三候選操作集合不合理。如果候選操作里有太多相似的操作比如同時(shí)有3x3卷積和5x5卷積搜索過(guò)程可能會(huì)在這兩個(gè)操作之間反復(fù)橫跳。建議候選操作集合保持多樣性但不要有太多冗余。第四batch size太小。搜索階段的batch size如果太小梯度的方差會(huì)很大導(dǎo)致搜索不穩(wěn)定。建議batch size至少設(shè)置為32如果顯存允許設(shè)置為64更好。5.2 離散化之后性能崩塌的原因分析離散化之后性能崩塌是DARTS最被詬病的問(wèn)題之一。原因主要有兩個(gè)一是連續(xù)松弛和離散選擇之間的gap二是搜索階段的過(guò)擬合。連續(xù)松弛的gap是指搜索階段最優(yōu)的連續(xù)架構(gòu)離散化之后不一定還是最優(yōu)的。因?yàn)檫B續(xù)架構(gòu)是所有候選操作的加權(quán)平均而離散架構(gòu)只保留權(quán)重最大的那個(gè)操作。如果權(quán)重分布比較均勻離散化之后的信息損失會(huì)很大。搜索階段的過(guò)擬合是指架構(gòu)參數(shù)在搜索過(guò)程中可能過(guò)擬合了驗(yàn)證集。因?yàn)樗阉麟A段只訓(xùn)練50個(gè)epoch網(wǎng)絡(luò)權(quán)重還沒(méi)有充分收斂架構(gòu)參數(shù)可能選擇了一些在短期訓(xùn)練中表現(xiàn)好、但長(zhǎng)期訓(xùn)練中表現(xiàn)差的操作。緩解這個(gè)問(wèn)題的方法有幾個(gè)。第一增加搜索階段的epoch數(shù)讓網(wǎng)絡(luò)權(quán)重更充分地收斂。第二在搜索階段加入正則化比如dropout或者權(quán)重衰減。第三在離散化之后對(duì)架構(gòu)做微調(diào)比如調(diào)整某些邊的操作看看性能是否提升。第四使用一些改進(jìn)的DARTS變體比如PC-DARTS、P-DARTS等這些方法在搜索穩(wěn)定性和離散化性能上都有改進(jìn)。5.3 顯存不足的排查與優(yōu)化方案顯存不足是DARTS實(shí)操中最常見(jiàn)的問(wèn)題。排查思路是這樣的首先看搜索階段的batch size是不是太大如果是先減小batch size。然后看候選操作數(shù)量是不是太多如果是減少候選操作。然后看是不是用了二階近似如果是改成一階近似。最后看是不是可以用混合精度訓(xùn)練。如果以上方法都不行還可以考慮梯度累積。梯度累積的思路是用小的batch size計(jì)算梯度累積多次之后再更新參數(shù)。這樣可以在不增加顯存占用的情況下模擬大的batch size。PyTorch里實(shí)現(xiàn)梯度累積很簡(jiǎn)單就是在backward之后不立即step而是累積幾次之后再step。for i, (input, target) in enumerate(train_loader): output model(input) loss criterion(output, target) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()這段代碼里accumulation_steps是累積次數(shù)。需要注意的是loss要除以accumulation_steps否則梯度會(huì)放大。5.4 常見(jiàn)問(wèn)題速查表問(wèn)題現(xiàn)象可能原因排查方法解決方案搜索階段loss震蕩架構(gòu)參數(shù)學(xué)習(xí)率太大打印alpha的梯度范數(shù)降低學(xué)習(xí)率到1e-4搜索階段loss不下降學(xué)習(xí)率太小或數(shù)據(jù)劃分不合理檢查驗(yàn)證集大小和分布增大學(xué)習(xí)率或調(diào)整數(shù)據(jù)劃分離散化后性能崩塌連續(xù)松弛gap或搜索過(guò)擬合對(duì)比搜索和評(píng)估階段的精度增加搜索epoch或使用改進(jìn)變體顯存不足batch size太大或候選操作太多用nvidia-smi監(jiān)控顯存減小batch size或減少候選操作評(píng)估階段精度波動(dòng)大隨機(jī)種子影響用多個(gè)隨機(jī)種子跑取平均值或增加訓(xùn)練epoch訓(xùn)練速度慢數(shù)據(jù)加載瓶頸或模型太大用profiler分析增加數(shù)據(jù)加載線程或減小模型6. 從DARTS出發(fā)的擴(kuò)展與改進(jìn)方向6.1 搜索效率的進(jìn)一步優(yōu)化DARTS雖然比強(qiáng)化學(xué)習(xí)的方法快很多但搜索階段仍然需要幾個(gè)GPU天。后續(xù)有很多工作試圖進(jìn)一步優(yōu)化搜索效率。比如PC-DARTS通過(guò)部分通道連接來(lái)減少顯存占用和計(jì)算量P-DARTS通過(guò)逐步增加網(wǎng)絡(luò)深度來(lái)穩(wěn)定搜索過(guò)程Fair-DARTS通過(guò)公平比較不同操作來(lái)避免搜索偏差。這些改進(jìn)的核心思路大同小異要么減少搜索階段的計(jì)算量要么提高搜索的穩(wěn)定性。PC-DARTS的部分通道連接是一個(gè)很實(shí)用的技巧它把每個(gè)操作的輸入通道分成兩部分一部分用于搜索一部分用于保留信息。這樣可以在不損失性能的情況下把顯存占用降低一半左右。6.2 搜索空間的擴(kuò)展與遷移DARTS的搜索空間最初是為圖像分類設(shè)計(jì)的后來(lái)被擴(kuò)展到其他任務(wù)比如語(yǔ)義分割、目標(biāo)檢測(cè)、語(yǔ)音識(shí)別等。擴(kuò)展搜索空間的關(guān)鍵是定義合適的候選操作和cell結(jié)構(gòu)。比如在語(yǔ)義分割里需要加入空洞卷積和上采樣操作在目標(biāo)檢測(cè)里需要加入多尺度特征融合的操作。遷移到其他數(shù)據(jù)集時(shí)需要注意數(shù)據(jù)集的規(guī)模和復(fù)雜度。如果數(shù)據(jù)集比較小搜索階段容易過(guò)擬合需要加強(qiáng)正則化如果數(shù)據(jù)集比較大搜索階段需要更多的epoch來(lái)充分收斂。6.3 與其他技術(shù)的結(jié)合DARTS可以和很多其他技術(shù)結(jié)合比如知識(shí)蒸餾、量化、剪枝等。知識(shí)蒸餾可以用一個(gè)大模型來(lái)指導(dǎo)搜索過(guò)程讓搜索出來(lái)的架構(gòu)更好地模仿大模型的行為。量化可以在搜索階段就考慮量化誤差搜索出對(duì)量化更友好的架構(gòu)。剪枝可以在搜索完成后對(duì)架構(gòu)做進(jìn)一步的壓縮。這些結(jié)合方式的核心思想是把DARTS作為一個(gè)架構(gòu)搜索的工具和其他技術(shù)協(xié)同使用而不是孤立地看待。在實(shí)際業(yè)務(wù)場(chǎng)景里往往需要綜合考慮精度、速度、模型大小等多個(gè)指標(biāo)單一的技術(shù)很難滿足所有需求。7. 我個(gè)人在實(shí)際操作中的幾點(diǎn)體會(huì)跑DARTS這段時(shí)間最大的感受是論文里的公式很優(yōu)雅但代碼里的細(xì)節(jié)很磨人。很多在論文里一筆帶過(guò)的東西在實(shí)際操作中都是坑。比如數(shù)據(jù)劃分的比例、學(xué)習(xí)率的調(diào)度、候選操作的選擇這些在論文里可能只有一句話但在代碼里需要反復(fù)調(diào)試。另外一個(gè)體會(huì)是不要迷信官方代碼的結(jié)果。官方代碼在CIFAR-10上能跑到97%以上的精度但換到其他數(shù)據(jù)集或者換個(gè)隨機(jī)種子結(jié)果可能差很多。所以復(fù)現(xiàn)的時(shí)候一定要多跑幾次取平均值不要被單次結(jié)果迷惑。最后DARTS的思想比具體的實(shí)現(xiàn)更重要。連續(xù)松弛、雙層優(yōu)化、cell結(jié)構(gòu)這些思想可以遷移到很多其他問(wèn)題里。即使你不做NAS了解這些思想對(duì)理解其他優(yōu)化問(wèn)題也有幫助。 DARTS實(shí)戰(zhàn)：從論文到代碼的神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索避坑指南
1. 為什么值得花時(shí)間啃透DARTS如果你正在做神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索NAS方向的研究或者工程落地大概率繞不開(kāi)DARTS這篇工作。全稱Differentiable Architecture Search核心思路一句話就能說(shuō)清楚把離散的架構(gòu)選擇松弛成連續(xù)的、可微的權(quán)重然后用梯度下降同時(shí)優(yōu)化網(wǎng)絡(luò)權(quán)重和架構(gòu)參數(shù)。這個(gè)想法在2018年提出來(lái)的時(shí)候直接炸了NAS這個(gè)圈子因?yàn)樗言拘枰獛浊PU天的搜索過(guò)程壓縮到了幾個(gè)GPU天。我自己第一次讀這篇論文的時(shí)候感覺(jué)公式推導(dǎo)很順代碼應(yīng)該也不難。結(jié)果真正把官方代碼拉下來(lái)跑通、再自己改結(jié)構(gòu)的時(shí)候才發(fā)現(xiàn)坑遠(yuǎn)比想象中多。從搜索階段的顯存爆炸到離散化之后的性能崩塌再到評(píng)估階段的復(fù)現(xiàn)差異每一步都有大量論文里不會(huì)寫(xiě)的細(xì)節(jié)。這篇文章就是把我從讀論文到跑代碼、再到自己改結(jié)構(gòu)這一整條鏈路上積累的東西全部攤開(kāi)講包括核心原理的拆解、代碼關(guān)鍵模塊的逐行分析、實(shí)操中遇到的典型問(wèn)題和排查思路。適合的讀者有三類一是剛接觸NAS、想通過(guò)DARTS入門(mén)的研究生或者算法工程師二是已經(jīng)跑過(guò)官方代碼但結(jié)果對(duì)不上、想搞清楚哪里出了問(wèn)題的同學(xué)三是想把DARTS的思想遷移到自己業(yè)務(wù)場(chǎng)景里的從業(yè)者。不管你是哪一類我都會(huì)盡量把“為什么這么做”講透而不是只貼一段代碼讓你抄。2. DARTS的核心設(shè)計(jì)思路拆解2.1 從離散搜索到連續(xù)松弛這個(gè)彎是怎么拐過(guò)來(lái)的傳統(tǒng)的NAS方法比如基于強(qiáng)化學(xué)習(xí)或者進(jìn)化算法的方案本質(zhì)上是在一個(gè)離散的搜索空間里做黑盒優(yōu)化。搜索空間里每個(gè)候選架構(gòu)都是一個(gè)獨(dú)立的點(diǎn)你需要完整訓(xùn)練它、評(píng)估它才能知道它好不好。這就導(dǎo)致搜索成本極高因?yàn)槊看卧u(píng)估都意味著一次完整的訓(xùn)練過(guò)程。DARTS的關(guān)鍵洞察在于我不需要把每個(gè)候選架構(gòu)當(dāng)成獨(dú)立的點(diǎn)我可以把搜索空間表示成一張有向無(wú)環(huán)圖每條邊上的操作是一個(gè)混合操作混合的權(quán)重就是架構(gòu)參數(shù)。這樣一來(lái)架構(gòu)搜索就變成了對(duì)這些連續(xù)權(quán)重的優(yōu)化問(wèn)題可以直接用梯度下降來(lái)解。具體來(lái)說(shuō)假設(shè)一個(gè)單元cell里有若干節(jié)點(diǎn)節(jié)點(diǎn)之間的邊上有若干候選操作比如3x3卷積、5x5卷積、跳連、池化等。DARTS的做法是在每條邊上放一個(gè)softmax加權(quán)的混合操作混合操作 sum over o of (exp(alpha_o) / sum over o of exp(alpha_o)) * o(x)其中alpha就是架構(gòu)參數(shù)。搜索的過(guò)程就是同時(shí)優(yōu)化網(wǎng)絡(luò)權(quán)重w和架構(gòu)參數(shù)alpha。訓(xùn)練完成后每條邊保留softmax權(quán)重最大的那個(gè)操作得到一個(gè)離散的架構(gòu)。這個(gè)松弛的好處是顯而易見(jiàn)的搜索過(guò)程變成了可微的可以用標(biāo)準(zhǔn)的反向傳播來(lái)優(yōu)化效率提升了好幾個(gè)數(shù)量級(jí)。但代價(jià)也很明顯松弛之后的連續(xù)空間和離散空間之間存在gap搜索時(shí)最優(yōu)的連續(xù)架構(gòu)離散化之后不一定還是最優(yōu)的。這個(gè)問(wèn)題在后面會(huì)詳細(xì)展開(kāi)。2.2 雙層優(yōu)化為什么不能簡(jiǎn)單地把a(bǔ)lpha和w一起訓(xùn)DARTS的優(yōu)化目標(biāo)是一個(gè)雙層優(yōu)化問(wèn)題bilevel optimization。外層優(yōu)化架構(gòu)參數(shù)alpha目標(biāo)是讓在驗(yàn)證集上的損失最小內(nèi)層優(yōu)化網(wǎng)絡(luò)權(quán)重w目標(biāo)是讓在訓(xùn)練集上的損失最小。用公式表示就是min over alpha of L_val(w*(alpha), alpha) s.t. w*(alpha) argmin over w of L_train(w, alpha)為什么不能直接把a(bǔ)lpha和w放在一起用訓(xùn)練集損失來(lái)優(yōu)化因?yàn)槟菢訒?huì)導(dǎo)致架構(gòu)參數(shù)過(guò)擬合訓(xùn)練集搜索出來(lái)的架構(gòu)在驗(yàn)證集上表現(xiàn)很差。雙層優(yōu)化的設(shè)計(jì)就是為了讓架構(gòu)參數(shù)在驗(yàn)證集上做選擇保證泛化性。但雙層優(yōu)化有個(gè)麻煩w*(alpha)本身是內(nèi)層優(yōu)化的結(jié)果要對(duì)alpha求導(dǎo)需要用到隱函數(shù)定理或者近似。DARTS采用的是一種一階近似的方法通過(guò)一步梯度下降來(lái)近似w*(alpha)然后對(duì)alpha求導(dǎo)。這個(gè)近似在論文里有詳細(xì)的推導(dǎo)代碼里也有對(duì)應(yīng)的實(shí)現(xiàn)。實(shí)際操作中這個(gè)近似會(huì)帶來(lái)一些問(wèn)題。比如搜索階段的batch size不能太小否則梯度的方差會(huì)很大導(dǎo)致alpha的優(yōu)化不穩(wěn)定。另外訓(xùn)練集和驗(yàn)證集的劃分也很關(guān)鍵如果驗(yàn)證集太小或者分布和訓(xùn)練集差異太大搜索出來(lái)的架構(gòu)可能會(huì)偏向于某些特定的操作。2.3 搜索空間設(shè)計(jì)cell結(jié)構(gòu)的選擇與權(quán)衡DARTS的搜索空間是基于cell的。一個(gè)cell是一個(gè)有向無(wú)環(huán)圖包含若干個(gè)節(jié)點(diǎn)和邊。DARTS定義了兩種cellnormal cell和reduction cell。Normal cell保持特征圖的空間分辨率不變r(jià)eduction cell將空間分辨率減半。每個(gè)cell有兩個(gè)輸入節(jié)點(diǎn)分別是前兩層的輸出和一個(gè)輸出節(jié)點(diǎn)。中間的節(jié)點(diǎn)通過(guò)邊連接每條邊上的操作從預(yù)定義的候選操作集合中選擇。候選操作通常包括3x3深度可分離卷積、5x5深度可分離卷積、3x3空洞卷積、5x5空洞卷積、3x3最大池化、3x3平均池化、identity跳連、zero無(wú)連接。這個(gè)設(shè)計(jì)的好處是搜索空間相對(duì)緊湊搜索出來(lái)的cell可以堆疊成任意深度的網(wǎng)絡(luò)。但也有一些限制比如cell的結(jié)構(gòu)是固定的只有邊上的操作是可搜索的節(jié)點(diǎn)的連接方式也是預(yù)定義的不能搜索。這些限制在一定程度上約束了搜索空間的表達(dá)能力但換來(lái)了搜索效率的大幅提升。我在實(shí)際使用中發(fā)現(xiàn)候選操作集合的選擇對(duì)搜索結(jié)果影響很大。如果池化操作太多搜索出來(lái)的架構(gòu)會(huì)傾向于大量使用池化導(dǎo)致模型容量不足如果卷積操作太多搜索又會(huì)偏向于堆疊卷積參數(shù)量爆炸。通常的做法是保持一個(gè)相對(duì)平衡的候選集合并根據(jù)具體任務(wù)做微調(diào)。3. 代碼關(guān)鍵模塊逐行解析3.1 搜索空間的構(gòu)建從代碼角度看cell是怎么定義的官方代碼里搜索空間的構(gòu)建主要在model_search.py這個(gè)文件里。核心類是Cell它接收一個(gè)steps參數(shù)表示cell里有多少個(gè)中間節(jié)點(diǎn)和一個(gè)multiplier參數(shù)表示輸出節(jié)點(diǎn)由多少個(gè)中間節(jié)點(diǎn)拼接而成。每個(gè)cell的構(gòu)建過(guò)程大致是這樣的首先定義兩個(gè)輸入節(jié)點(diǎn)前兩層的輸出然后對(duì)于每個(gè)中間節(jié)點(diǎn)從前面的所有節(jié)點(diǎn)中選擇兩個(gè)作為輸入在這兩條邊上分別放置混合操作?；旌喜僮鞯臋?quán)重就是架構(gòu)參數(shù)alpha。代碼里有一個(gè)關(guān)鍵的數(shù)據(jù)結(jié)構(gòu)self._ops它是一個(gè)ModuleList每個(gè)元素是一個(gè)MixedOp對(duì)應(yīng)一條邊上的混合操作。MixedOp內(nèi)部維護(hù)了一個(gè)ModuleList包含所有候選操作以及一個(gè)alpha參數(shù)用于計(jì)算softmax權(quán)重。class MixedOp(nn.Module): def __init__(self, C, stride): super(MixedOp, self).__init__() self._ops nn.ModuleList() for primitive in PRIMITIVES: op OPS[primitive](C, stride, False) self._ops.append(op) def forward(self, x, weights): return sum(w * op(x) for w, op in zip(weights, self._ops))這段代碼看起來(lái)很簡(jiǎn)單但有幾個(gè)細(xì)節(jié)值得注意。第一weights是通過(guò)對(duì)alpha做softmax得到的這個(gè)softmax是在forward里計(jì)算的而不是在__init__里。第二sum這里用的是Python內(nèi)置的sum而不是torch.sum這意味著每個(gè)操作的輸出是逐個(gè)相加的顯存占用會(huì)比較高。第三OPS是一個(gè)字典把操作名映射到具體的nn.Module這個(gè)字典的定義在文件開(kāi)頭。在實(shí)際使用中我發(fā)現(xiàn)MixedOp的顯存占用是一個(gè)大問(wèn)題。因?yàn)槊織l邊上都要計(jì)算所有候選操作的輸出然后加權(quán)求和所以顯存占用是候選操作數(shù)量的倍數(shù)。如果候選操作有8個(gè)那顯存占用就是普通網(wǎng)絡(luò)的8倍左右。這也是為什么DARTS搜索階段需要比較大的顯存。3.2 架構(gòu)參數(shù)的初始化與更新alpha的優(yōu)化細(xì)節(jié)架構(gòu)參數(shù)alpha的初始化方式對(duì)搜索結(jié)果有一定影響。官方代碼里alpha是用torch.nn.Parameter定義的初始化為0。這意味著初始時(shí)所有操作的softmax權(quán)重是均勻的每條邊上的混合操作是所有候選操作的等權(quán)平均。alpha的更新是通過(guò)優(yōu)化器來(lái)完成的。官方代碼里用的是torch.optim.Adam學(xué)習(xí)率設(shè)置為3e-4權(quán)重衰減設(shè)置為1e-3。這個(gè)學(xué)習(xí)率比網(wǎng)絡(luò)權(quán)重的學(xué)習(xí)率通常0.025小很多原因是架構(gòu)參數(shù)需要更精細(xì)的調(diào)整學(xué)習(xí)率太大會(huì)導(dǎo)致搜索過(guò)程震蕩。在train_search.py里有一個(gè)關(guān)鍵的步驟是計(jì)算架構(gòu)參數(shù)的梯度。由于雙層優(yōu)化的存在這個(gè)梯度不能直接用驗(yàn)證集損失對(duì)alpha求導(dǎo)得到而是需要先對(duì)網(wǎng)絡(luò)權(quán)重做一步虛擬更新然后再計(jì)算梯度。代碼里對(duì)應(yīng)的邏輯是# 計(jì)算虛擬更新后的網(wǎng)絡(luò)權(quán)重 unrolled_model self._compute_unrolled_model(input, target, eta, network_optimizer) # 計(jì)算架構(gòu)參數(shù)的梯度 unrolled_loss self._loss(unrolled_model, self._steps, input_valid, target_valid) unrolled_loss.backward()這里的_compute_unrolled_model就是做一步虛擬更新eta是虛擬更新的學(xué)習(xí)率。這個(gè)步驟是DARTS的核心也是計(jì)算量最大的部分因?yàn)樗枰Ａ粲?jì)算圖顯存占用會(huì)進(jìn)一步增加。我在實(shí)際使用中發(fā)現(xiàn)如果顯存不夠可以把unrolled設(shè)置為False改用一階近似。這樣雖然搜索精度會(huì)有所下降但顯存占用會(huì)大幅降低。對(duì)于顯存有限的場(chǎng)景這是一個(gè)很實(shí)用的折中方案。3.3 離散化與架構(gòu)導(dǎo)出從連續(xù)權(quán)重到最終網(wǎng)絡(luò)搜索完成后需要把連續(xù)權(quán)重離散化得到最終的架構(gòu)。這個(gè)過(guò)程在model_search.py的genotype方法里實(shí)現(xiàn)。核心邏輯是對(duì)于每條邊選擇softmax權(quán)重最大的操作對(duì)于每個(gè)中間節(jié)點(diǎn)選擇權(quán)重最大的兩條邊作為輸入。def genotype(self): def _parse(weights): gene [] n 2 start 0 for i in range(self._steps): end start n W weights[start:end].copy() edges sorted(range(i 2), keylambda x: -max(W[x][k] for k in range(len(W[x])) if k ! PRIMITIVES.index(none)))[:2] for j in edges: k_best None for k in range(len(W[j])): if k ! PRIMITIVES.index(none): if k_best is None or W[j][k] W[j][k_best]: k_best k gene.append((PRIMITIVES[k_best], j)) start end n 1 return gene # ...這段代碼有幾個(gè)細(xì)節(jié)需要注意。第一none操作被排除在外也就是說(shuō)如果某條邊的none操作權(quán)重最大這條邊會(huì)被忽略。第二每個(gè)中間節(jié)點(diǎn)選擇兩條輸入邊這是DARTS的默認(rèn)設(shè)置。第三gene是一個(gè)列表每個(gè)元素是一個(gè)(操作名, 輸入節(jié)點(diǎn)索引)的元組。離散化之后通常會(huì)重新訓(xùn)練一個(gè)網(wǎng)絡(luò)來(lái)評(píng)估架構(gòu)的性能。這個(gè)重新訓(xùn)練的過(guò)程和普通網(wǎng)絡(luò)訓(xùn)練類似但有一些細(xì)節(jié)需要注意比如學(xué)習(xí)率調(diào)度、數(shù)據(jù)增強(qiáng)、正則化等。官方代碼里用的是標(biāo)準(zhǔn)的訓(xùn)練流程但在實(shí)際使用中這些超參數(shù)對(duì)最終性能影響很大。4. 實(shí)操過(guò)程中的關(guān)鍵環(huán)節(jié)與避坑指南4.1 環(huán)境配置與數(shù)據(jù)準(zhǔn)備從零開(kāi)始跑通官方代碼先把環(huán)境搭起來(lái)。官方代碼基于PyTorch建議用1.0以上的版本但不要用太新的版本因?yàn)橛行〢PI可能不兼容。Python版本建議3.6到3.8之間。CUDA版本根據(jù)你的顯卡來(lái)選但要注意PyTorch版本和CUDA版本的對(duì)應(yīng)關(guān)系。數(shù)據(jù)準(zhǔn)備方面官方代碼支持CIFAR-10和ImageNet。CIFAR-10可以直接用torchvision.datasets.CIFAR10下載ImageNet需要自己準(zhǔn)備。數(shù)據(jù)增強(qiáng)方面官方代碼用了標(biāo)準(zhǔn)的隨機(jī)裁剪和隨機(jī)翻轉(zhuǎn)另外還用了cutout。Cutout的實(shí)現(xiàn)很簡(jiǎn)單就是在圖像上隨機(jī)挖一個(gè)16x16的方塊把它置零。這個(gè)操作對(duì)防止過(guò)擬合很有效。class Cutout(object): def __init__(self, length): self.length length def __call__(self, img): h, w img.size(1), img.size(2) mask np.ones((h, w), np.float32) y np.random.randint(h) x np.random.randint(w) y1 np.clip(y - self.length // 2, 0, h) y2 np.clip(y self.length // 2, 0, h) x1 np.clip(x - self.length // 2, 0, w) x2 np.clip(x self.length // 2, 0, w) mask[y1: y2, x1: x2] 0. img torch.from_numpy(img.numpy() * mask) return img這段代碼看起來(lái)簡(jiǎn)單但有一個(gè)坑img.numpy()要求img是CPU上的tensor如果數(shù)據(jù)加載器返回的是CUDA tensor這里會(huì)報(bào)錯(cuò)。所以cutout要在數(shù)據(jù)加載階段做而不是在訓(xùn)練循環(huán)里做。4.2 搜索階段的參數(shù)設(shè)置與顯存優(yōu)化搜索階段的參數(shù)設(shè)置很關(guān)鍵。官方代碼里搜索階段訓(xùn)練50個(gè)epochbatch size設(shè)置為64網(wǎng)絡(luò)權(quán)重的學(xué)習(xí)率是0.025架構(gòu)參數(shù)的學(xué)習(xí)率是3e-4。這些參數(shù)在CIFAR-10上表現(xiàn)不錯(cuò)但換到其他數(shù)據(jù)集可能需要調(diào)整。顯存優(yōu)化是搜索階段的一個(gè)大問(wèn)題。因?yàn)槊織l邊上都要計(jì)算所有候選操作的輸出顯存占用是普通網(wǎng)絡(luò)的數(shù)倍。如果顯存不夠可以嘗試以下幾個(gè)方案減小batch size。這是最直接的方法但會(huì)影響搜索穩(wěn)定性。減少候選操作數(shù)量。比如去掉一些不常用的操作只保留核心的幾個(gè)。使用一階近似。把unrolled設(shè)置為False可以大幅降低顯存占用。使用混合精度訓(xùn)練。PyTorch的amp模塊可以自動(dòng)做混合精度顯存占用可以減少30%到50%。我自己的經(jīng)驗(yàn)是如果顯存只有8GB建議用一階近似加混合精度batch size設(shè)置為32左右。如果顯存有16GB以上可以用二階近似batch size設(shè)置為64。4.3 評(píng)估階段的重新訓(xùn)練與性能復(fù)現(xiàn)搜索完成后會(huì)得到一個(gè)離散的架構(gòu)。這個(gè)架構(gòu)需要重新訓(xùn)練來(lái)評(píng)估性能。官方代碼里評(píng)估階段訓(xùn)練600個(gè)epochbatch size設(shè)置為96學(xué)習(xí)率用余弦退火調(diào)度初始學(xué)習(xí)率0.025。評(píng)估階段有幾個(gè)坑需要注意。第一數(shù)據(jù)增強(qiáng)的強(qiáng)度要夠否則容易過(guò)擬合。除了標(biāo)準(zhǔn)的隨機(jī)裁剪和翻轉(zhuǎn)還可以加cutout、mixup、cutmix等。第二學(xué)習(xí)率調(diào)度很關(guān)鍵余弦退火比階梯下降效果更好。第三權(quán)重衰減要設(shè)置合理太大或太小都會(huì)影響性能。我在復(fù)現(xiàn)官方結(jié)果的時(shí)候發(fā)現(xiàn)評(píng)估階段的性能波動(dòng)很大。同樣的架構(gòu)不同的隨機(jī)種子最終精度可能差0.5%到1%。所以如果要對(duì)比不同架構(gòu)的性能最好用多個(gè)隨機(jī)種子跑取平均值。另外官方代碼里的評(píng)估階段用的是CIFAR-10的完整訓(xùn)練集但搜索階段用的是訓(xùn)練集的50%作為訓(xùn)練集另外50%作為驗(yàn)證集。這個(gè)劃分方式在評(píng)估階段要改回來(lái)否則會(huì)浪費(fèi)一半的訓(xùn)練數(shù)據(jù)。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 搜索階段loss不下降或者震蕩怎么辦搜索階段loss不下降或者震蕩通常有幾個(gè)原因。第一學(xué)習(xí)率設(shè)置不當(dāng)。架構(gòu)參數(shù)的學(xué)習(xí)率如果太大會(huì)導(dǎo)致搜索過(guò)程震蕩如果太小搜索會(huì)停滯。建議從3e-4開(kāi)始試如果震蕩就降到1e-4如果停滯就升到1e-3。第二數(shù)據(jù)劃分不合理。如果驗(yàn)證集太小或者分布和訓(xùn)練集差異太大架構(gòu)參數(shù)的梯度會(huì)很不穩(wěn)定。建議驗(yàn)證集至少占訓(xùn)練集的30%到50%。第三候選操作集合不合理。如果候選操作里有太多相似的操作比如同時(shí)有3x3卷積和5x5卷積搜索過(guò)程可能會(huì)在這兩個(gè)操作之間反復(fù)橫跳。建議候選操作集合保持多樣性但不要有太多冗余。第四batch size太小。搜索階段的batch size如果太小梯度的方差會(huì)很大導(dǎo)致搜索不穩(wěn)定。建議batch size至少設(shè)置為32如果顯存允許設(shè)置為64更好。5.2 離散化之后性能崩塌的原因分析離散化之后性能崩塌是DARTS最被詬病的問(wèn)題之一。原因主要有兩個(gè)一是連續(xù)松弛和離散選擇之間的gap二是搜索階段的過(guò)擬合。連續(xù)松弛的gap是指搜索階段最優(yōu)的連續(xù)架構(gòu)離散化之后不一定還是最優(yōu)的。因?yàn)檫B續(xù)架構(gòu)是所有候選操作的加權(quán)平均而離散架構(gòu)只保留權(quán)重最大的那個(gè)操作。如果權(quán)重分布比較均勻離散化之后的信息損失會(huì)很大。搜索階段的過(guò)擬合是指架構(gòu)參數(shù)在搜索過(guò)程中可能過(guò)擬合了驗(yàn)證集。因?yàn)樗阉麟A段只訓(xùn)練50個(gè)epoch網(wǎng)絡(luò)權(quán)重還沒(méi)有充分收斂架構(gòu)參數(shù)可能選擇了一些在短期訓(xùn)練中表現(xiàn)好、但長(zhǎng)期訓(xùn)練中表現(xiàn)差的操作。緩解這個(gè)問(wèn)題的方法有幾個(gè)。第一增加搜索階段的epoch數(shù)讓網(wǎng)絡(luò)權(quán)重更充分地收斂。第二在搜索階段加入正則化比如dropout或者權(quán)重衰減。第三在離散化之后對(duì)架構(gòu)做微調(diào)比如調(diào)整某些邊的操作看看性能是否提升。第四使用一些改進(jìn)的DARTS變體比如PC-DARTS、P-DARTS等這些方法在搜索穩(wěn)定性和離散化性能上都有改進(jìn)。5.3 顯存不足的排查與優(yōu)化方案顯存不足是DARTS實(shí)操中最常見(jiàn)的問(wèn)題。排查思路是這樣的首先看搜索階段的batch size是不是太大如果是先減小batch size。然后看候選操作數(shù)量是不是太多如果是減少候選操作。然后看是不是用了二階近似如果是改成一階近似。最后看是不是可以用混合精度訓(xùn)練。如果以上方法都不行還可以考慮梯度累積。梯度累積的思路是用小的batch size計(jì)算梯度累積多次之后再更新參數(shù)。這樣可以在不增加顯存占用的情況下模擬大的batch size。PyTorch里實(shí)現(xiàn)梯度累積很簡(jiǎn)單就是在backward之后不立即step而是累積幾次之后再step。for i, (input, target) in enumerate(train_loader): output model(input) loss criterion(output, target) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()這段代碼里accumulation_steps是累積次數(shù)。需要注意的是loss要除以accumulation_steps否則梯度會(huì)放大。5.4 常見(jiàn)問(wèn)題速查表問(wèn)題現(xiàn)象可能原因排查方法解決方案搜索階段loss震蕩架構(gòu)參數(shù)學(xué)習(xí)率太大打印alpha的梯度范數(shù)降低學(xué)習(xí)率到1e-4搜索階段loss不下降學(xué)習(xí)率太小或數(shù)據(jù)劃分不合理檢查驗(yàn)證集大小和分布增大學(xué)習(xí)率或調(diào)整數(shù)據(jù)劃分離散化后性能崩塌連續(xù)松弛gap或搜索過(guò)擬合對(duì)比搜索和評(píng)估階段的精度增加搜索epoch或使用改進(jìn)變體顯存不足batch size太大或候選操作太多用nvidia-smi監(jiān)控顯存減小batch size或減少候選操作評(píng)估階段精度波動(dòng)大隨機(jī)種子影響用多個(gè)隨機(jī)種子跑取平均值或增加訓(xùn)練epoch訓(xùn)練速度慢數(shù)據(jù)加載瓶頸或模型太大用profiler分析增加數(shù)據(jù)加載線程或減小模型6. 從DARTS出發(fā)的擴(kuò)展與改進(jìn)方向6.1 搜索效率的進(jìn)一步優(yōu)化DARTS雖然比強(qiáng)化學(xué)習(xí)的方法快很多但搜索階段仍然需要幾個(gè)GPU天。后續(xù)有很多工作試圖進(jìn)一步優(yōu)化搜索效率。比如PC-DARTS通過(guò)部分通道連接來(lái)減少顯存占用和計(jì)算量P-DARTS通過(guò)逐步增加網(wǎng)絡(luò)深度來(lái)穩(wěn)定搜索過(guò)程Fair-DARTS通過(guò)公平比較不同操作來(lái)避免搜索偏差。這些改進(jìn)的核心思路大同小異要么減少搜索階段的計(jì)算量要么提高搜索的穩(wěn)定性。PC-DARTS的部分通道連接是一個(gè)很實(shí)用的技巧它把每個(gè)操作的輸入通道分成兩部分一部分用于搜索一部分用于保留信息。這樣可以在不損失性能的情況下把顯存占用降低一半左右。6.2 搜索空間的擴(kuò)展與遷移DARTS的搜索空間最初是為圖像分類設(shè)計(jì)的后來(lái)被擴(kuò)展到其他任務(wù)比如語(yǔ)義分割、目標(biāo)檢測(cè)、語(yǔ)音識(shí)別等。擴(kuò)展搜索空間的關(guān)鍵是定義合適的候選操作和cell結(jié)構(gòu)。比如在語(yǔ)義分割里需要加入空洞卷積和上采樣操作在目標(biāo)檢測(cè)里需要加入多尺度特征融合的操作。遷移到其他數(shù)據(jù)集時(shí)需要注意數(shù)據(jù)集的規(guī)模和復(fù)雜度。如果數(shù)據(jù)集比較小搜索階段容易過(guò)擬合需要加強(qiáng)正則化如果數(shù)據(jù)集比較大搜索階段需要更多的epoch來(lái)充分收斂。6.3 與其他技術(shù)的結(jié)合DARTS可以和很多其他技術(shù)結(jié)合比如知識(shí)蒸餾、量化、剪枝等。知識(shí)蒸餾可以用一個(gè)大模型來(lái)指導(dǎo)搜索過(guò)程讓搜索出來(lái)的架構(gòu)更好地模仿大模型的行為。量化可以在搜索階段就考慮量化誤差搜索出對(duì)量化更友好的架構(gòu)。剪枝可以在搜索完成后對(duì)架構(gòu)做進(jìn)一步的壓縮。這些結(jié)合方式的核心思想是把DARTS作為一個(gè)架構(gòu)搜索的工具和其他技術(shù)協(xié)同使用而不是孤立地看待。在實(shí)際業(yè)務(wù)場(chǎng)景里往往需要綜合考慮精度、速度、模型大小等多個(gè)指標(biāo)單一的技術(shù)很難滿足所有需求。7. 我個(gè)人在實(shí)際操作中的幾點(diǎn)體會(huì)跑DARTS這段時(shí)間最大的感受是論文里的公式很優(yōu)雅但代碼里的細(xì)節(jié)很磨人。很多在論文里一筆帶過(guò)的東西在實(shí)際操作中都是坑。比如數(shù)據(jù)劃分的比例、學(xué)習(xí)率的調(diào)度、候選操作的選擇這些在論文里可能只有一句話但在代碼里需要反復(fù)調(diào)試。另外一個(gè)體會(huì)是不要迷信官方代碼的結(jié)果。官方代碼在CIFAR-10上能跑到97%以上的精度但換到其他數(shù)據(jù)集或者換個(gè)隨機(jī)種子結(jié)果可能差很多。所以復(fù)現(xiàn)的時(shí)候一定要多跑幾次取平均值不要被單次結(jié)果迷惑。最后DARTS的思想比具體的實(shí)現(xiàn)更重要。連續(xù)松弛、雙層優(yōu)化、cell結(jié)構(gòu)這些思想可以遷移到很多其他問(wèn)題里。即使你不做NAS了解這些思想對(duì)理解其他優(yōu)化問(wèn)題也有幫助。