YOLOv5における凍結レイヤーを用いた転移学習

📚 This guide explains how to freeze YOLOv5 🚀 layers when implementing transfer learning. Transfer learning is a powerful machine learning (ML) technique that allows you to quickly retrain a model on new data without retraining the entire network from scratch. By freezing the weights of initial layers and only updating the parameters of later layers, you can significantly reduce computational resource requirements and training time. However, this approach might slightly impact the final model accuracy.

始める前に

まず、YOLOv5リポジトリをクローンし、requirements.txtに記載されている必要な依存関係をインストールします。Python>=3.8.0環境とPyTorch>=1.8がインストールされていることを確認してください。事前学習済みモデルと必要なデータセットは、最新のYOLOv5リリースから自動的にダウンロードされます。

git clone https://github.com/ultralytics/yolov5 # clone repository
cd yolov5
pip install -r requirements.txt # install dependencies

レイヤー凍結の仕組み

ニューラルネットワークのレイヤーを凍結すると、そのパラメータ(重みとバイアス)が学習プロセス中に更新されるのを防ぐことができます。PyTorchでは、レイヤーのテンソルのrequires_grad属性をFalseに設定することでこれを実現します。その結果、逆伝播中にこれらのレイヤーの勾配が計算されず、計算量とメモリが節約されます。

YOLOv5がその学習スクリプトにおいてどのようにレイヤーの凍結を実装しているかを以下に示します。

# Freeze specified layers
freeze = [f"model.{x}." for x in range(freeze)]  # Define layers to freeze based on module index
for k, v in model.named_parameters():
    v.requires_grad = True  # Ensure all parameters are initially trainable
    if any(x in k for x in freeze):
        print(f"Freezing layer: {k}")
        v.requires_grad = False  # Disable gradient calculation for frozen layers

モデルアーキテクチャの調査

どのレイヤーを凍結するかを決定するには、YOLOv5モデルの構造を理解することが不可欠です。以下のPythonスニペットを使用して、すべてのモジュールとそのパラメータの名前を確認できます。

# Assuming 'model' is your loaded YOLOv5 model instance
for name, param in model.named_parameters():
    print(name)

"""
Example Output:
model.0.conv.conv.weight
model.0.conv.bn.weight
model.0.conv.bn.bias
model.1.conv.weight
model.1.bn.weight
model.1.bn.bias
model.2.cv1.conv.weight
model.2.cv1.bn.weight
...
"""

The YOLOv5 architecture typically consists of a backbone (layers 0-9 in standard configurations like YOLOv5s/m/l/x) responsible for feature extraction, and a head (the remaining layers) which performs object detection.

# Example YOLOv5 v6.0 backbone structure
backbone:
    # [from, number, module, args]
    - [-1, 1, Conv, [64, 6, 2, 2]]  # Layer 0: Initial convolution (P1/2 stride)
    - [-1, 1, Conv, [128, 3, 2]] # Layer 1: Downsampling convolution (P2/4 stride)
    - [-1, 3, C3, [128]]          # Layer 2: C3 module
    - [-1, 1, Conv, [256, 3, 2]] # Layer 3: Downsampling convolution (P3/8 stride)
    - [-1, 6, C3, [256]]          # Layer 4: C3 module
    - [-1, 1, Conv, [512, 3, 2]] # Layer 5: Downsampling convolution (P4/16 stride)
    - [-1, 9, C3, [512]]          # Layer 6: C3 module
    - [-1, 1, Conv, [1024, 3, 2]]# Layer 7: Downsampling convolution (P5/32 stride)
    - [-1, 3, C3, [1024]]         # Layer 8: C3 module
    - [-1, 1, SPPF, [1024, 5]]    # Layer 9: Spatial Pyramid Pooling Fast

# Example YOLOv5 v6.0 head structure
head:
    - [-1, 1, Conv, [512, 1, 1]] # Layer 10
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]] # Layer 11
    - [[-1, 6], 1, Concat, [1]] # Layer 12: Concatenate with backbone P4 (from layer 6)
    - [-1, 3, C3, [512, False]] # Layer 13: C3 module
    # ... subsequent head layers for feature fusion and detection

凍結オプション

学習コマンドの--freeze引数を使用して、どのレイヤーを凍結するかを制御できます。この引数は、凍結されない最初のモジュールのインデックスを指定します。このインデックスより前のすべてのモジュールの重みが凍結されます。どのインデックスが特定のブロックに対応するかを確認する必要がある場合は、model.modelnn.Sequential)を使用してモジュールの順序を確認してください。

バックボーンのみの凍結

COCOのような大規模なデータセットから学習した汎用的な特徴抽出能力を維持しつつ、モデルを新しいオブジェクトクラスに適応させる際によく用いられる、バックボーン全体(レイヤー0から9まで)を凍結する方法です。

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 10

この戦略は、ターゲットデータセットが元の学習データ(COCOなど)と類似した低レベルの視覚的特徴(エッジ、テクスチャ)を共有しているが、異なるオブジェクトカテゴリを含んでいる場合に有効です。

最終検出レイヤー以外をすべて凍結

ネットワーク全体をほぼ凍結し、最終的な出力畳み込みレイヤー(Detectモジュールの一部。通常は最後のモジュール、例:YOLOv5sではモジュール24)のみを学習可能にする場合:

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 24

このアプローチは、学習済みの大部分の特徴を保持したまま、主に出力クラスの数に合わせてモデルを調整する必要がある場合に便利です。ファインチューニングにおける計算リソースを最小限に抑えることができます。

パフォーマンスの比較

To illustrate the effects of freezing layers, we trained YOLOv5m on the Pascal VOC dataset for 50 epochs, starting from the official COCO pretrained weights (yolov5m.pt). We compared three scenarios: training all layers (--freeze 0), freezing the backbone (--freeze 10), and freezing all but the final detection layers (--freeze 24).

# Example command for training with backbone frozen
python train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp hyp.finetune.yaml --freeze 10

精度結果

結果から、レイヤーの凍結によって学習時間を大幅に短縮できるものの、最終的なmAP (mean Average Precision)がわずかに低下する可能性があることがわかりました。通常、すべてのレイヤーを学習させる方が精度は最も高くなりますが、より多くのレイヤーを凍結すると、パフォーマンスがわずかに低下する代わりに学習が高速化します。

異なる凍結戦略を比較した学習時のmAP50結果 学習中のmAP50比較

異なる凍結戦略を比較した学習時のmAP50-95結果 学習中のmAP50-95比較

YOLOv5 frozen layer training performance *Summary table of performance metrics*

リソース利用状況

より多くのレイヤーを凍結することで、GPUのメモリ要件と全体的な利用率が大幅に削減されます。これにより、ハードウェアリソースが限られている環境で転移学習を行う際に魅力的な選択肢となり、凍結しない場合よりも大きなモデルの学習や、より大きな画像サイズの使用が可能になります。

学習中のGPUメモリ割り当て割合 GPUメモリ割り当て (%)

学習中のGPUメモリ利用率 GPU利用率 (%)

レイヤー凍結の活用場面

転移学習中のレイヤー凍結は、以下のような状況で特に有利です。

  1. 計算リソースが限られている場合: GPUメモリや処理能力に制限があるとき。
  2. 小規模なデータセット: ターゲットデータセットが元の事前学習データセットよりも大幅に小さい場合、凍結は過学習を防ぐのに役立ちます。
  3. 迅速なプロトタイピング: 既存のモデルを新しいタスクやドメインに迅速に適応させて初期評価を行う必要があるとき。
  4. 類似した特徴ドメイン: 新しいデータセットの低レベル特徴が、モデルが事前学習されたデータセットと非常に類似している場合。

転移学習のニュアンスについては、用語集エントリで詳細を確認してください。また、パフォーマンスを最適化するためにハイパーパラメータチューニングなどの手法も検討してください。

サポートされている環境

Ultralyticsでは、CUDACuDNNPythonPyTorchなどの主要な依存関係がプリインストールされた、すぐに使えるさまざまな環境を提供しています。

プロジェクトステータス

YOLOv5 Continuous Integration Status

このバッジは、すべてのYOLOv5 GitHub Actions継続的インテグレーション(CI)テストが正常に通過していることを示しています。これらのCIテストは、学習検証推論エクスポートベンチマークといった主要な操作全般にわたり、YOLOv5の機能とパフォーマンスを厳密に評価します。これにより、macOS、Windows、Ubuntu上で一貫した信頼性の高い動作が保証され、24時間ごとおよび新しいコードコミットごとに自動的に実行されます。

コメント