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