记录一次 mmclassification 自定义数据训练和推理

总体参考如下(还有其他CSDN和知乎贴子):

欢迎来到 MMClassification 中文教程! — MMClassification 0.21.0 文档

mmclassification/getting_started.md at master · open-mmlab/mmclassification · GitHub


1. 环境安装

 除了安装基础的python,pytorch等,重点是mmcv和mmcls!

 由于要用到开发场景,不要用pip安装封装好的包,用官方建议(官方install那一步也有讲):

pip install openmim

mim install -e .


2. 代码

直接clone的mmlab官方源码:

GitHub - open-mmlab/mmclassification at devicon-default.png?t=M276https://github.com/open-mmlab/mmclassification/tree/dev


3. 数据集

自己数据集有个文件夹,假设我是猫狗数据集dogs_cats,下边有如下几个文件夹:

 记录一次 mmclassification 自定义数据训练和推理_第1张图片

以猫狗数据集为例,其中假设你已经在train、test和val三个文件夹中放好了图片了(注意看我红线部分,他们三个里边都分别有两个文件夹,分别是class1和class2, class2装的猫图片)。

记录一次 mmclassification 自定义数据训练和推理_第2张图片

 meta中应该有如下几个文件:

记录一次 mmclassification 自定义数据训练和推理_第3张图片

 其中先看classmap.txt,其中内容写好类名和他的类代号:

记录一次 mmclassification 自定义数据训练和推理_第4张图片

 train.txt,test.txt, val.txt用一个脚本 data_deal.py 生成,运行完了,他们内容下边这样:

记录一次 mmclassification 自定义数据训练和推理_第5张图片

至于 data_deal.py脚本参考:

# -*- coding:utf8 -*-

import os
import glob
import re

'''
生成train.txt  val.txt  test.txt
'''

# 需要改为您自己的路径
root_dir = "/xxx/data/cats_dogs"
# 在该路径下有train val test meta三个文件夹
train_dir = os.path.join(root_dir, "train")
val_dir = os.path.join(root_dir, "val")
test_dir = os.path.join(root_dir, "test")
meta_dir = os.path.join(root_dir, "meta")

def generate_txt(images_dir,map_dict):
    # 读取所有文件名
    imgs_dirs = glob.glob(images_dir+"/*/*")

    imgs_dirs = [ii.replace('\\', '/') for ii in imgs_dirs]
    images_dir = images_dir.replace('\\', '/')

    typename = images_dir.split("/")[-1]
    target_txt_path = os.path.join(meta_dir, typename+".txt")
    f = open(target_txt_path, "w")
    # 遍历所有图片名
    for img_dir in imgs_dirs:
        # 获取第一级目录名称
        filename = img_dir.split("/")[-2]
        num = map_dict[filename]
        # 写入文件
        relate_name = re.findall(typename+"/([\w / - .]*)",img_dir)
        f.write(relate_name[0]+" "+num+"\n")

def get_map_dict():
    # 读取所有类别映射关系
    class_map_dict = {}
    with open(os.path.join(meta_dir, "classmap.txt"),"r") as F:
        lines = F.readlines()
        for line in lines:
            line = line.split("\n")[0]
            filename, cls, num = line.split(" ")
            class_map_dict[filename] = num
    return class_map_dict

if __name__ == '__main__':

    class_map_dict = get_map_dict()
    generate_txt(images_dir=train_dir, map_dict=class_map_dict)
    generate_txt(images_dir=val_dir, map_dict=class_map_dict)
    generate_txt(images_dir=test_dir, map_dict=class_map_dict)

4. 配置文件

   注:以下描述的路径均在项目根目录的基础路径下!

(1)新建 mmcls\datasets\a_mydataset.py ,自己数据集的定义脚本

        (主要是CLASSES = ["dog", "cat"]这一句):

# -*- coding:utf8 -*-

import numpy as np

from .builder import DATASETS
from .base_dataset import BaseDataset

@DATASETS.register_module()
class MyDataset(BaseDataset):
    CLASSES = ["dog", "cat"]
    def load_annotations(self):
        assert isinstance(self.ann_file, str)
        data_infos = []
        with open(self.ann_file) as f:
            samples = [x.strip().split(' ') for x in f.readlines()]
            for filename, gt_label in samples:
                info = {'img_prefix': self.data_prefix}
                info['img_info'] = {'filename': filename}
                info['gt_label'] = np.array(gt_label, dtype=np.int64)
                data_infos.append(info)
            return data_infos

(2)mmcls\datasets\__init__.py   需要对自己的数据在init里添加,进行注册:

记录一次 mmclassification 自定义数据训练和推理_第6张图片

 (3)新建 config/_base_/datasets/a_mydataset.py  (主要参考 imagenet_bs32.py 对应改为自己的数据路径)

# dataset settings
dataset_type = 'MyDataset'
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)  # coco数据训练时原始的预处理方式
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', size=224),  # 随机缩放裁剪
    dict(type='RandomFlip', flip_prob=0.5, direction='horizontal'),  # 随机翻转
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='ToTensor', keys=['gt_label']),
    dict(type='Collect', keys=['img', 'gt_label'])  # 图像和标签的集合
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', size=(256, -1)),
    dict(type='CenterCrop', crop_size=224),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='ImageToTensor', keys=['img']),
    dict(type='Collect', keys=['img'])
]

data_root = '/xxx/data/dogs_cats'  # 你的数据集根目录
data = dict(
    samples_per_gpu=5,  # dataloader.batch_size == self.samples_per_gpu  # 每批次样本数量
    workers_per_gpu=2,  # dataloader.num_workers == self.workers_per_gpu  # 1的话表示只有一个进程加载数据
    train=dict(
        type=dataset_type,
        data_prefix=data_root + '/train',
        ann_file=data_root + '/meta/train.txt',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_prefix=data_root + '/val',
        ann_file=data_root + '/meta/val.txt',
        pipeline=test_pipeline),
    test=dict(
        # replace `data/val` with `data/test` for standard test
        type=dataset_type,
        data_prefix=data_root + '/val',
        ann_file=data_root + '/meta/val.txt',
        pipeline=test_pipeline))
evaluation = dict(interval=1, metric='accuracy')  # 分类,用accuracy

(4)新建 config/_base_/model/a_resnet18_mydataset.py, 主要是参考resnet18.py

         (网络的设置:注意是head里边的num_classes设置成自己的类别数量)

# model settings
model = dict(
    type='ImageClassifier',
    backbone=dict(  # backbone: RES-NET
        type='ResNet',
        depth=18,  # resnet18
        num_stages=4,  # resnet18 --- block-depth分别是(2,2,2,2)
        out_indices=(3, ),  # 取第三阶段,也就是最后一阶段的输出,resnet最后输出特征纬度是512
        style='pytorch',
        # torchvision默认预训练模型,这里默认下载放到/home/yons/.cache/torch/hub/checkpoints
        # init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet18'),
        # 预训练模型也可以用官方应用imagenet训练好的
        # https://github.com/open-mmlab/mmclassification/blob/dev/docs/en/model_zoo.md 可下载checkpoint和查看更多,这里下载对应 resnet18_8xb32_in1k.py 的预训练设置
        init_cfg=dict(
            type='Pretrained',
            checkpoint='/xxx/checkpoints/'resnet18_batch256_imagenet_20200708-34ab8f90.pth',  
            prefix='backbone',
        ),
        frozen_stages=2,  # 加入预训练模型时候,冻结前两层(实验证明,冻结部分权重,效果更好)
        
    ),
    neck=dict(type='GlobalAveragePooling'),  # neck:平均池化
    head=dict(
        type='LinearClsHead',
        num_classes=2,  # 二分类(因为是猫狗数据集了,不是imagenet)
        in_channels=512,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),  # 分类,交叉熵
        topk=(1, 5),
        # topk=(1, 2)  # so you would have to check output.shape and make sure dim1 is larger or equal to maxk.
    )
)

(5)新建 config/_base_/schedules/a_imagenet_bs256.py

          或者直接用imagenet_bs256.py,     定义优化器,学习率,和迭代次数等等。

# optimizer
optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)  # momentum:学习率动量 weight_decay:权重惩罚,正则化
optimizer_config = dict(grad_clip=None)  # 设置梯度截断(裁剪)阈值,防止梯度无限大爆炸
# learning policy
lr_config = dict(policy='step', step=[30, 60, 90])  # 100次迭代,每隔设定的这几次,lr就降低0.1倍
runner = dict(type='EpochBasedRunner', max_epochs=100)

(6) config/_base_/default_runtime.py,  设置多少批次打印日志,多少次迭代保存一次模型等等。

# checkpoint saving
checkpoint_config = dict(interval=20)  # 多少次迭代保存一次模型
# yapf:disable
log_config = dict(
    # interval=100,
    interval=40,  # 多少批次 打印一次
    hooks=[
        dict(type='TextLoggerHook'),
        # dict(type='TensorboardLoggerHook')
    ])
# yapf:enable

dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]

(7) 新建 config/resnet/a_resnet18_b32x8_mydataset.py (这里就是模型配置主脚本了,借鉴resnet18_b32x8_imagenet的训练配置,按照如下顺序定义自己的配置文件。

# b32指的batch-size,8指gpu
_base_ = [
    '../_base_/models/a_resnet18_mydataset.py',  # model config
    '../_base_/datasets/a_mydataset.py',  # data config
    '../_base_/schedules/a_imagenet_bs256.py',  # schedules config
    '../_base_/default_runtime.py'  # runtime config
]

5. 训练

tools/train.py 中修改config和work-dir  ,然后运行:python tools/train.py

记录一次 mmclassification 自定义数据训练和推理_第7张图片

 训练完了过后,不出意外的话,训练结果会保存到设定的train_result/0310下(这个可以自己随便设置):

记录一次 mmclassification 自定义数据训练和推理_第8张图片


6. 推理验证

(1)一张图评估, 就新建 test_one.py  用下边:

# Copyright (c) OpenMMLab. All rights reserved.

from argparse import ArgumentParser
import argparse
from mmcls.apis import inference_model, init_model, show_result_pyplot


def main():

    # ### 测试的args,用于debug ###################
    args = argparse.Namespace()

    # == 本地 ===========
    args.img = r'\xxx\dogs_cats\val\class2\cat.5658.jpg'
    args.config = 'configs/resnet/a_resnet18_b32x8_mydataset.py'
    args.checkpoint = 'train_result/0310/epoch_100.pth'
    args.device = 'cpu'

    # build the model from a config file and a checkpoint file
    model = init_model(args.config, args.checkpoint, device=args.device)
    # test a single image
    result = inference_model(model, args.img)
    # show the results
    show_result_pyplot(model, args.img, result)


if __name__ == '__main__':
    main()

   效果还不错!

记录一次 mmclassification 自定义数据训练和推理_第9张图片

(2)对val.txt集中测试, 参考tools/test.py, 只需要对应改一下,再 python tools/test.py 即可!

记录一次 mmclassification 自定义数据训练和推理_第10张图片


7. 小结

  (1) 环境一定要装好,mmcls要安装开发版的,不是应用版的;

(2)关于数据可以用你自己的,我网上下载的猫狗二分类图片;

(3)所有步骤,一环扣一环,不可忽略了;

(3)有问题可以到mmlab的官方论坛知乎,或者github中查看,没问题的,加油!奥利给~

你可能感兴趣的:(计算机视觉,人工智能,深度学习,pytorch,深度学习,人工智能)