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