跳轉到主要內容
部落格

如何使用 TorchVision 最新基元訓練最先進的模型

作者: 2021 年 11 月 18 日2024 年 11 月 15 日無評論

幾周前,TorchVision v0.11 釋出,其中包含眾多新的基元、模型和訓練方案改進,從而實現了最先進 (SOTA) 的結果。該專案被命名為“內建電池的 TorchVision”,旨在使我們的庫現代化。我們希望透過使用常見的構建塊,使研究人員能夠更輕鬆地復現論文並進行研究。此外,我們渴望為應用機器學習從業者提供必要的工具,使其能夠使用與研究中相同的 SOTA 技術來訓練自己的資料模型。最後,我們希望重新整理我們預訓練的權重,併為使用者提供更好的開箱即用模型,希望他們能構建更好的應用程式。

儘管還有很多工作要做,但我們想與您分享上述工作的一些令人興奮的成果。我們將展示如何使用 TorchVision 中包含的新工具,在 ResNet50 等競爭激烈且經過充分研究的架構上獲得最先進的結果 [1]。我們將分享用於將基線提高超過 4.7 個準確點,達到 80.9% 最終 top-1 準確率的精確方案,並分享推導新訓練過程的旅程。此外,我們將展示該方案很好地推廣到其他模型變體和系列。我們希望上述內容能影響未來研究,以開發更強大的通用訓練方法,並激勵社群採納併為我們的努力做出貢獻。

結果

使用我們在 ResNet50 上發現的新訓練方案,我們重新整理了以下模型的預訓練權重

模型 準確率@1 準確率@5
ResNet50 80.858 95.434
ResNet101 81.886 95.780
ResNet152 82.284 96.002
ResNeXt50-32x4d 81.198 95.340

請注意,除了 ResNet50 之外,所有模型的準確性都可以透過稍微調整其訓練引數來進一步提高,但我們的重點是擁有一套適用於所有模型的單一穩健方案。

更新:我們已經重新整理了 TorchVision 中大多數流行的分類模型,您可以在此部落格文章中找到詳細資訊。

目前有兩種方法可以使用模型的最新權重。

使用多預訓練權重 API

我們目前正在開發一種新的原型機制,該機制將擴充套件 TorchVision 的模型構建器方法以支援多個權重。除了權重之外,我們還儲存有用的元資料(例如標籤、準確性、方案連結等)以及使用模型所需的預處理轉換。示例

  from PIL import Image
  from torchvision import prototype as P
  img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")
   
  # Initialize model
  weights = P.models.ResNet50_Weights.IMAGENET1K_V2
  model = P.models.resnet50(weights=weights)
  model.eval()

  # Initialize inference transforms
  preprocess = weights.transforms()
   
  # Apply inference preprocessing transforms
  batch = preprocess(img).unsqueeze(0)
  prediction = model(batch).squeeze(0).softmax(0)
   
  # Make predictions
  label = prediction.argmax().item()
  score = prediction[label].item()
   
  # Use meta to get the labels
  category_name = weights.meta['categories'][label]
  print(f"{category_name}: {100 * score}%")

使用傳統 API

不希望使用原型 API 的使用者可以選擇透過以下方法訪問傳統 API 中的新權重

  from torchvision.models import resnet
   
  # Overwrite the URL of the previous weights
  resnet.model_urls["resnet50"] = "https://download.pytorch.org/models/resnet50-11ad3fa6.pth"
   
  # Initialize the model using the legacy API
  model = resnet.resnet50(pretrained=True)
   
  # TODO: Apply preprocessing + call the model
  # ...

訓練方案

我們的目標是使用 TorchVision 新引入的基元來推匯出一個強大的新訓練方案,該方案在 ImageNet 上從頭開始訓練,不使用額外的外部資料時,能為香草 ResNet50 架構實現最先進的結果。儘管透過使用特定於架構的技巧 [2] 可以進一步提高準確性,但我們決定不包含它們,以便該方案可以用於其他架構。我們的方案[3][4][5][6][7] 嚴重側重於簡潔性,並建立在 FAIR 的工作基礎上。我們的發現與 Wightman 等人 [7] 的並行研究一致,他們也報告了透過關注訓練方案而實現的主要準確性改進。

閒話少說,以下是我們方案的主要引數

  # Optimizer & LR scheme
  ngpus=8,
  batch_size=128,  # per GPU

  epochs=600, 
  opt='sgd',  
  momentum=0.9,

  lr=0.5, 
  lr_scheduler='cosineannealinglr', 
  lr_warmup_epochs=5, 
  lr_warmup_method='linear', 
  lr_warmup_decay=0.01, 


  # Regularization and Augmentation
  weight_decay=2e-05, 
  norm_weight_decay=0.0,

  label_smoothing=0.1, 
  mixup_alpha=0.2, 
  cutmix_alpha=1.0, 
  auto_augment='ta_wide', 
  random_erase=0.1, 
  
  ra_sampler=True,
  ra_reps=4,


  # EMA configuration
  model_ema=True, 
  model_ema_steps=32, 
  model_ema_decay=0.99998, 


  # Resizing
  interpolation='bilinear', 
  val_resize_size=232, 
  val_crop_size=224, 
  train_crop_size=176,

使用我們的標準訓練參考指令碼,我們可以使用以下命令訓練 ResNet50

torchrun --nproc_per_node=8 train.py --model resnet50 --batch-size 128 --lr 0.5 \
--lr-scheduler cosineannealinglr --lr-warmup-epochs 5 --lr-warmup-method linear \
--auto-augment ta_wide --epochs 600 --random-erase 0.1 --weight-decay 0.00002 \
--norm-weight-decay 0.0 --label-smoothing 0.1 --mixup-alpha 0.2 --cutmix-alpha 1.0 \
--train-crop-size 176 --model-ema --val-resize-size 232 --ra-sampler --ra-reps 4

方法論

在我們的探索過程中,我們牢記以下幾個原則

  1. 訓練是一個隨機過程,我們試圖最佳化的驗證指標是一個隨機變數。這是由於所採用的隨機權重初始化方案和訓練過程中隨機效應的存在。這意味著我們不能透過單次執行來評估方案更改的效果。標準做法是進行多次執行(通常為 3 到 5 次)並研究彙總統計資料(例如平均值、標準差、中位數、最大值等)。
  2. 不同引數之間通常存在顯著的相互作用,特別是對於側重於正則化和減少過擬合的技術。因此,更改其中一個引數的值可能會影響其他引數的最佳配置。為了解決這個問題,可以採用貪婪搜尋方法(通常會導致次優結果但可追蹤的實驗)或應用網格搜尋(會導致更好的結果但計算成本高)。在這項工作中,我們混合使用了這兩種方法。
  3. 非確定性或引入噪聲的技術通常需要更長的訓練週期才能提高模型效能。為了使事情易於處理,我們最初使用短訓練週期(少量 epoch)來決定哪些路徑可以提前消除,哪些應該使用更長的訓練週期進行探索。
  4. 由於重複實驗,存在驗證資料集過擬合的風險 [8]。為了減輕部分風險,我們只應用提供顯著準確性改進的訓練最佳化,並使用 K 折交叉驗證來驗證在驗證集上進行的最佳化。此外,我們確認我們的方案成分在其他模型上也能很好地推廣,而我們並未最佳化這些模型的超引數。

關鍵準確性改進細分

正如早期部落格文章中所討論的,訓練模型並非一個準確性單調遞增的旅程,該過程涉及大量的回溯。為了量化每次最佳化的效果,下面我們嘗試展示一個從 TorchVision 原始方案開始推導最終方案的理想化線性旅程。我們想澄清的是,這過度簡化了我們實際遵循的路徑,因此應慎重對待。

Cumulative Accuracy Improvements for ResNet50

在下表中,我們總結了在基線之上疊加增量改進的效能。除非另有說明,否則我們報告 3 次執行中 Acc@1 最佳的模型

準確率@1 準確率@5 增量差異 絕對差異
ResNet50 基線 76.130 92.862 0.000 0.000
+ LR 最佳化 76.494 93.198 0.364 0.364
+ TrivialAugment 76.806 93.272 0.312 0.676
+ 長時間訓練 78.606 94.052 1.800 2.476
+ 隨機擦除 78.796 94.094 0.190 2.666
+ 標籤平滑 79.114 94.374 0.318 2.984
+ Mixup 79.232 94.536 0.118 3.102
+ Cutmix 79.510 94.642 0.278 3.380
+ 權重衰減調整 80.036 94.746 0.526 3.906
+ FixRes 緩解 80.196 94.672 0.160 4.066
+ EMA 80.450 94.908 0.254 4.320
+ 推理大小調整 * 80.674 95.166 0.224 4.544
+ 重複增強 ** 80.858 95.434 0.184 4.728

*推理大小的調整是在最後一個模型之上完成的。詳情見下文。

** 文章釋出後完成的社群貢獻。詳情見下文。

基線

我們的基線是 TorchVision 之前釋出的 ResNet50 模型。它使用以下方案進行訓練

  # Optimizer & LR scheme
  ngpus=8,
  batch_size=32,  # per GPU

  epochs=90, 
  opt='sgd',  
  momentum=0.9,

  lr=0.1, 
  lr_scheduler='steplr', 
  lr_step_size=30, 
  lr_gamma=0.1, 


  # Regularization
  weight_decay=1e-4,


  # Resizing
  interpolation='bilinear', 
  val_resize_size=256, 
  val_crop_size=224, 
  train_crop_size=224,

上述大多數引數都是我們訓練指令碼中的預設值。我們將在此基線之上逐步引入最佳化,直到最終達到最終方案。

LR 最佳化

我們可以應用一些引數更新來提高訓練的準確性和速度。這可以透過增加批次大小和調整學習率來實現。另一種常見的方法是應用預熱並逐漸增加學習率。這在學習率非常高時特別有利,有助於訓練在早期 epoch 的穩定性。最後,另一個最佳化是應用餘弦排程器在 epoch 期間調整學習率。餘弦排程器的一個巨大優勢是無需最佳化超引數,這減少了我們的搜尋空間。

以下是在基線方案之上應用的額外最佳化。請注意,我們進行了多次實驗以確定引數的最佳配置

  batch_size=128,  # per GPU

  lr=0.5, 
  lr_scheduler='cosineannealinglr', 
  lr_warmup_epochs=5, 
  lr_warmup_method='linear', 
  lr_warmup_decay=0.01,

與基線相比,上述最佳化將我們的 top-1 準確率提高了 0.364 個點。請注意,為了結合不同的學習率策略,我們使用了新引入的 SequentialLR 排程器。

TrivialAugment

原始模型使用基本的增強轉換進行訓練,例如隨機裁剪和水平翻轉。提高準確率的一種簡單方法是應用更復雜的“自動增強”技術。對我們來說表現最好的是 TrivialAugment [9],它非常簡單,可以認為是“無引數”的,這意味著它可以幫助我們進一步縮小搜尋空間。

以下是在上一步基礎上應用的更新

auto_augment='ta_wide',

與上一步相比,TrivialAugment 的使用使我們的 top-1 準確率提高了 0.312 個點。

長時間訓練

當我們的方案包含隨機行為的成分時,更長的訓練週期是有益的。更具體地說,當我們開始新增越來越多的引入噪聲的技術時,增加 epoch 數量變得至關重要。請注意,在我們探索的早期階段,我們使用了大約 200 個 epoch 的相對較短的週期,隨著我們開始縮小大部分引數範圍,後來增加到 400 個 epoch,最後在方案的最終版本中增加到 600 個 epoch。

下面我們看到在早期步驟基礎上應用的更新

epochs=600,

這使我們的 top-1 準確率在之前的基礎上進一步提高了 1.8 個點。這是我們在此迭代過程中觀察到的最大增幅。值得注意的是,這種單一最佳化的效果被誇大且有些誤導。僅僅在舊基線之上增加 epoch 數量並不會產生如此顯著的改進。然而,學習率最佳化與強大增強策略的結合有助於模型從更長的週期中受益。還值得一提的是,我們如此早地引入長時間訓練週期的原因是在接下來的步驟中我們將引入需要更多 epoch 才能提供良好結果的技術。

隨機擦除

另一種已知有助於分類準確性的資料增強技術是隨機擦除 [10][11]。它通常與自動增強方法結合使用,由於其正則化效果,通常會進一步提高準確性。在我們的實驗中,我們僅透過網格搜尋調整了應用該方法的機率,發現將其機率保持在較低水平(通常約為 10%)是有益的。

這是在之前基礎上引入的額外引數

random_erase=0.1,

應用隨機擦除使我們的 Acc@1 進一步提高了 0.190 個點。

標籤平滑

減少過擬合的一個好技術是阻止模型過於自信。這可以透過使用標籤平滑 [12] 軟化真實標籤來實現。需要指定一個控制平滑程度(越高越強)的單一引數。儘管可以透過網格搜尋對其進行最佳化,但我們發現 0.05-0.15 左右的值會產生相似的結果,因此為了避免過擬合,我們使用了引入它的論文中相同的值。

下面我們可以找到在此步驟中新增的額外配置

label_smoothing=0.1,

我們使用了 PyTorch 新引入的 CrossEntropyLoss 標籤平滑引數,這使我們的準確率額外提高了 0.318 個點。

Mixup 和 Cutmix

Mixup 和 Cutmix [13][14] 是兩種常用於產生 SOTA 結果的資料增強技術。它們都透過不僅軟化標籤而且軟化影像來提供強大的正則化效果。在我們的設定中,我們發現以相同的機率隨機應用其中一種是有益的。每種技術都由一個超引數 alpha 進行引數化,該超引數控制 Beta 分佈的形狀,從中取樣平滑機率。我們進行了非常有限的網格搜尋,主要關注論文中提出的常見值。

以下是兩種技術 alpha 引數的最佳值

mixup_alpha=0.2, 
cutmix_alpha=1.0,

應用 Mixup 使我們的準確率提高了 0.118 個點,將其與 Cutmix 結合使用可額外提高 0.278 個點。

權重衰減調整

我們的標準方案使用 L2 正則化來減少過擬合。權重衰減引數控制正則化的程度(越大越強),並且預設情況下普遍應用於模型的所有學習引數。在此方案中,我們對標準方法應用了兩項最佳化。首先,我們執行網格搜尋以調整權重衰減引數,其次,我們對歸一化層的引數停用權重衰減。

下面您可以找到我們方案中權重衰減的最佳配置

weight_decay=2e-05, 
norm_weight_decay=0.0,

上述更新使我們的準確率進一步提高了 0.526 個點,為調整權重衰減對模型效能具有顯著影響這一已知事實提供了額外的實驗證據。我們將歸一化引數與其餘引數分離的方法受到 ClassyVision 方法的啟發。

FixRes 緩解

在我們實驗早期發現的一個重要特性是,如果驗證期間使用的解析度從訓練的 224×224 增加,模型的效能會顯著提高。這一效應在 FixRes 論文 [5] 中詳細研究,並提出了兩種緩解措施:a) 可以嘗試降低訓練解析度,以使驗證解析度上的準確性最大化;b) 可以透過兩階段訓練對模型進行微調,使其適應目標解析度。由於我們不想引入兩階段訓練,因此我們選擇了選項 a)。這意味著我們減小了訓練裁剪大小,從 224 開始,並使用網格搜尋找到在 224×224 解析度上最大化驗證的那個。

下面您可以看到我們的方案中使用的最佳值

val_crop_size=224, 
train_crop_size=176,

上述最佳化使我們的準確率額外提高了 0.160 個點,並將訓練速度提高了 10%。

值得注意的是,FixRes 效應仍然存在,這意味著當我們增加解析度時,模型在驗證上表現仍然更好。此外,進一步縮小訓練裁剪大小實際上會損害準確性。這在直觀上是合理的,因為人們只能在關鍵細節開始從圖片中消失之前,將解析度降低到一定程度。最後,我們應該注意到,上述 FixRes 緩解措施似乎對與 ResNet50 深度相似的模型有益。具有更大感受野的更深變體似乎略微受到負面影響(通常為 0.1-0.2 個點)。因此,我們將這部分方案視為可選。下面我們可視化了使用 176 和 224 解析度訓練的最佳可用檢查點(包含完整方案)的效能

Best ResNet50 trained with 176 Resolution
Best ResNet50 trained with 224 Resolution

指數移動平均 (EMA)

EMA 是一種在不增加模型複雜性或推理時間的情況下提高模型準確性的技術。它對模型權重執行指數移動平均,這會提高準確性並使模型更穩定。平均每隔幾次迭代進行一次,其衰減引數透過網格搜尋進行調整。

下面您可以看到我們方案的最佳值

model_ema=True, 
model_ema_steps=32, 
model_ema_decay=0.99998,

與上一步相比,使用 EMA 將我們的準確率提高了 0.254 個點。請注意,TorchVision 的 EMA 實現 構建在 PyTorch 的 AveragedModel 類之上,其主要區別在於它不僅平均模型引數,還平均其緩衝區。此外,我們採用了 Pycls 的技巧,這使我們能夠以不依賴於 epoch 數量的方式引數化衰減。

推理調整大小

與過程中涉及使用不同引數訓練模型的所有其他步驟不同,此最佳化是在最終模型之上完成的。在推理期間,影像被調整為特定解析度,然後從中獲取一箇中央 224×224 裁剪。原始方案使用 256 的調整大小,這導致了與 FixRes 論文 [5] 中描述的類似差異。透過將此調整大小值更接近目標推理解析度,可以提高準確性。為了選擇該值,我們對區間 [224, 256] 進行了短網格搜尋,步長為 8。為了避免過擬合,該值使用一半驗證集選擇,並使用另一半進行確認。

下面您可以看到我們的方案中使用的最佳值

val_resize_size=232,

上述最佳化使我們的準確率提高了 0.224 個點。值得注意的是,ResNet50 的最佳值也適用於 ResNet101、ResNet152 和 ResNeXt50,這表明它可以在模型之間推廣

ResNet50 Inference Resize
ResNet101 Inference Resize
Best ResNet50 trained with 224 Resolution

[更新] 重複增強

重複增強 [15][16] 是另一種可以提高整體準確性的技術,已被其他強方案(如 [6][7])使用。社群貢獻者 Tal Ben-Nun 進一步改進 了我們的原始方案,建議使用 4 次重複訓練模型。他的貢獻是在本文釋出之後完成的。

下面您可以看到我們的方案中使用的最佳值

ra_sampler=True,
ra_reps=4,

上述是最終的最佳化,使我們的準確率提高了 0.184 個點。

已測試但未採用的最佳化

在研究的早期階段,我們嘗試了額外的技術、配置和最佳化。由於我們的目標是儘可能簡化方案,因此我們決定不包含任何未提供顯著改進的內容。以下是我們嘗試但未進入最終方案的一些方法

  • 最佳化器:使用更復雜的最佳化器,如 Adam、RMSProp 或帶有 Nesterov 動量的 SGD,並未比帶有動量的普通 SGD 提供顯著更好的結果。
  • 學習率排程器:我們嘗試了不同的學習率排程方案,例如 StepLR 和 Exponential。儘管後者傾向於與 EMA 更好地配合,但它通常需要額外的超引數,例如定義最小學習率才能很好地工作。相反,我們只使用餘弦退火將學習率衰減到零,並選擇具有最高準確率的檢查點。
  • 自動增強:我們嘗試了不同的增強策略,例如 AutoAugment 和 RandAugment。這些都沒有超越更簡單的無引數 TrivialAugment。
  • 插值:使用雙三次或最近鄰插值並未比雙線性插值提供顯著更好的結果。
  • 歸一化層:使用同步批次歸一化並未比使用常規批次歸一化產生顯著更好的結果。

鳴謝

我們衷心感謝 Piotr Dollar、Mannat Singh 和 Hugo Touvron 在方案開發過程中提供了他們的見解和反饋,以及他們之前的研究工作為我們的方案奠定了基礎。他們的支援對於取得上述成果至關重要。此外,我們還要感謝 Prabhat Roy、Kai Zhang、Yiwen Song、Joel Schlosser、Ilqar Ramazanli、Francisco Massa、Mannat Singh、Xiaoliang Dai、Samuel Gabriel、Allen Goodman 和 Tal Ben-Nun 對“內建電池”專案的貢獻。

參考文獻

  1. Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. “深度殘差學習用於影像識別”。
  2. Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, Mu Li. “卷積神經網路影像分類的技巧包”
  3. Piotr Dollár, Mannat Singh, Ross Girshick. “快速準確的模型縮放”
  4. Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, Ross Girshick. “早期卷積幫助 Transformers 更好地觀察”
  5. Hugo Touvron, Andrea Vedaldi, Matthijs Douze, Hervé Jégou. “修復訓練-測試解析度差異”
  6. Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, Hervé Jégou. “透過注意力訓練資料高效影像 transformer 和蒸餾”
  7. Ross Wightman, Hugo Touvron, Hervé Jégou. “ResNet 捲土重來:timm 中改進的訓練過程”
  8. Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, Vaishaal Shankar. “ImageNet 分類器能否推廣到 ImageNet?”
  9. Samuel G. Müller, Frank Hutter. “TrivialAugment:免調優但最先進的資料增強”
  10. Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, Yi Yang. “隨機擦除資料增強”
  11. Terrance DeVries, Graham W. Taylor. “透過 Cutout 改進卷積神經網路的正則化”
  12. Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna. “重新思考計算機視覺的 Inception 架構”
  13. Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz. “mixup:超越經驗風險最小化”
  14. Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, Youngjoon Yoo. “CutMix:用於訓練具有區域性特徵的強大分類器的正則化策略”
  15. Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, Daniel Soudry. “增強你的批次:用更大的批次進行更好的訓練”
  16. Maxim Berman, Hervé Jégou, Andrea Vedaldi, Iasonas Kokkinos, Matthijs Douze. “Multigrain:用於類別和例項的統一影像嵌入”