【图像分类】MMPretrain训练ImageNet格式自定义数据集

1. 安装环境

1.1. Conda虚拟环境

conda create --name mmpretrain python=3.8 -y
 
conda activate mmpretrain

1.2. 安装PyTorch

pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117

1.3. 安装mmcv

pip install -U openmim
mim install mmengine
mim install "mmcv==2.0.0"

1.4. clone库并setup

git clone https://github.com/open-mmlab/mmpretrain.git -b dev-1.x
cd mmpretrain
pip install -v -e .

2. 数据准备

Data

|--meta

|--train

       |--class1 ...

|--val

       |--class1 ...

|--test

       |--class1 ...

 代码：生成meta文件

import os
 
# test文件夹对应路径
test_path = 'train'
test_txt = 'meta/train.txt'
dirlist = os.listdir(test_path)
with open(test_txt, 'w') as f:  # 如果filename不存在会自动创建， 'w'表示写数据，写之前会清空文件中的原有数据！
    for i in os.listdir(test_path):
        print(i)
        for ii in os.listdir(os.path.join(test_path, i)):
            print(i+'/'+ii+" "+i)
            f.writelines(i+'/'+ii+" "+i+"\n")
f.close()
print("success")

3. 模型训练

3.1. 修改数据config

configs/_base_/datasets/config.py

(1) 根据数据保存位置修改data_root的路径

(2) 添加classes类别信息, classes=classes

3.2. 修改网络类别数

3.3. 设置预训练模型

选择使用的方法，下载预训练模型

设置configs/_base_/default_runtime.py

在load_from中设置路径

3.4. Setup

python setup.py install

3.5. 模型训练

python tools/train.py config.py

4. 模型测试

python tools/test.py configs.py weights.pth

5. 计算网络参数量

python tools/analysis_tools/get_flops.py config.py