在ImageNet 1k數據集上訓練yolov5m-cls分類模型

ImageNet 1k數據集簡介

ImageNet是CV領域非常出名的數據集， 其中ISLVRC2012數據集是Large Scale Visual Recognition Challenge 2012所用的數據集，包括：

• 訓練數據集ILSVRC2012_img_test.tar，裏面包括1000類共計1,281,167張圖片，大約138G
• 驗證數據集ILSVRC2012_img_val.tar，裏面包括1000類每類50張圖片，一共50000張圖片，大約6.3G

由於ISLVRC2012有1000類數據，所以很多論文把這個數據集叫做：ImageNet 1K。行業裏面漸漸約定俗成用這個數據集來測試模型結構，或者從零開始訓練一個全新的CNN主幹網絡(backbone)。

ImageNet的評價指標是固定的：top1 acc 和 top5 acc。基於ImageNet 1K訓練的模型，很容易跟已發表的模型比較，看看性能是否有提高

下載並解壓ImageNet 1k數據集

第一步，下載ILSVRC2012_img_test.tar和ILSVRC2012_img_val.tar

# parent
# ├── yolov5
# └── datasets
#     └── imagenet  ← downloads here

第二步，解壓ImageNet 1k數據集：創建解壓批處理文件：unpack_imagenet.sh

train=true
val=true
# Download/unzip train
if [ "$train" == "true" ]; then
  #wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar # download 138G, 1281167 images
  mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
  tar -xf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
  find . -name "*.tar" | while read NAME; do
    mkdir -p "${NAME%.tar}"
    tar -xf "${NAME}" -C "${NAME%.tar}"
    rm -f "${NAME}"
  done
  cd ..
fi

# Download/unzip val
if [ "$val" == "true" ]; then
  #wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar # download 6.3G, 50000 images
  mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xf ILSVRC2012_img_val.tar
  wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash # move into subdirs
fi

並在MINGW64中運行!

在ImageNet 1k數據集上訓練yolov5m-cls分類模型

克隆YOLOv5並安裝其依賴軟件包：

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

運行命令：

python classify/train.py --model yolov5m-cls.pt --data d:/datasets/imagenet --img 224 --batch-size 128 --workers 4

結果如下：

在訓練環境: Windows11 + i7-13700k + DDR5 6400MHz + 3060 12G 下實測最佳穩定訓練的參數爲：--img 224 --batch-size 128 --workers 4

DDR5 高頻內存 + PCIe 4.0的SSD，相對DDR4和HDD配置，對訓練速度提升作用非常大

單機單卡在imagenet 1k數據集上訓練yolov5m-cls模型，每輪(Epoch)耗時 36分鐘, 90輪共計54小時，對比4卡 A100的訓練時間10:06小時，預測8卡3060 12G訓練在imagenet 1k數據集上分類模型時間要遠低於4卡A100(該結論待找機器驗證)。

結論：感覺模型大了，例如，48.1M參數的YOLOv5x-cls模型，訓練速度就比yolov5m-cls明顯慢多了；大模型訓練，不僅考慮顯存大小，也要考慮顯卡的CUDA核心數量。基於ImageNet 1k數據集在RTX3060上訓練YOLOv5m-cls 20個Epoch，用時35.081小時，如下圖所示：