【mmpretrain/mmdeploy】mmpretrain的使用。将mmpretrain的模型转换为engine

使用mmpretrain中的resnet50来训练自己数据集的教程，我做的是两个类别的图像分类任务。

官网：https://mmpretrain.readthedocs.io/zh-cn/latest/index.html
参考：mmpretrain使用示例

我是直接在原来的resnet基础上进行修改的，其实应该复制一份再操作，这里图简单。

1. 打开configs/resnet/resnet18_8xb32_in1k.py

我们需要修改的就是这四个文件，分别找到对应的文件进行修改。

2.修改 mmpretrain/configs/base/models/resnet50.py

这里的type要查看一下文档。

3. 修改mmpretrain/configs/base/datasets/imagenet_bs32.py

# dataset settings
dataset_type = 'CustomDataset'
data_preprocessor = dict(
    num_classes=2,
    # RGB format normalization parameters
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    # convert image from BGR to RGB
    to_rgb=True,
)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='RandomResizedCrop', scale=224),
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),
    dict(type='PackInputs'),
]

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='ResizeEdge', scale=256, edge='short'),
    dict(type='CenterCrop', crop_size=224),
    dict(type='PackInputs'),
]

train_dataloader = dict(
    batch_size=32,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='data/gateDataset/train',
        #split='train',
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),
)

val_dataloader = dict(
    batch_size=32,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='data/gateDataset/val',
        #split='val',
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
)
val_evaluator = dict(type='Accuracy', topk=1 )


test_dataloader = dict(
    batch_size=32,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='data/gateDataset/test',
        
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
)
test_evaluator = dict(type='Accuracy', topk=1)

这里的test是我仿照前面的train和val自己写的。

这里的topk可以改成1，在我这个任务中。

对于他数据集：

这里的测试数据集也要按照训练的数据集格式来放置图片，不然有问题。

3. mmpretrain/configs/base/schedules/imagenet_bs256.py
   
   这个文件不用修改，这里主要是可以修改训练的epoch数。

4.mmpretrain/configs/base/default_runtime.py

# defaults to use registries in mmpretrain
default_scope = 'mmpretrain'

# configure default hooks
default_hooks = dict(
    # record the time of every iteration.
    timer=dict(type='IterTimerHook'),

    # print log every 100 iterations.
    logger=dict(type='LoggerHook', interval=100),

    # enable the parameter scheduler.
    param_scheduler=dict(type='ParamSchedulerHook'),

    # save checkpoint per epoch.
    checkpoint=dict(type='CheckpointHook', interval=1),

    # set sampler seed in distributed evrionment.
    sampler_seed=dict(type='DistSamplerSeedHook'),

    # validation results visualization, set True to enable it.
    visualization=dict(type='VisualizationHook', enable=False),
)

# configure environment
env_cfg = dict(
    # whether to enable cudnn benchmark
    cudnn_benchmark=False,

    # set multi process parameters
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),

    # set distributed parameters
    dist_cfg=dict(backend='nccl'),
)

# set visualizer
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(type='UniversalVisualizer', vis_backends=vis_backends)

# set log level
log_level = 'INFO'

# load from which checkpoint
load_from = None
#load_from = './resnet50_8xb32_in1k_20210831-ea4938fc.pth' 

# whether to resume training from the loaded checkpoint
resume = False

# Defaults to use random seed and disable `deterministic`
randomness = dict(seed=None, deterministic=False)

这个文件也不用修改，他有一个加载预训练模型的load_from ，如果有需要可以使用。

5. 训练的指令

python tools/train.py configs/resnet/resnet50_8xb32_in1k.py 

训练完成后，会出现一个work_dirs文件夹，模型都保存在里面。

6. 可视化训练过程。

python tools/analysis_tools/analyze_logs.py plot_curve work_dirs/resnet50_8xb32_in1k/20240807_052824/vis_data/20240807_052824.json  --keys loss accuracy/top1  --legend loss top1  --out results.jpg

python tools/analysis_tools/confusion_matrix.py configs/resnet/resnet50_8xb32_in1k.py work_dirs/resnet50_8xb32_in1k/epoch_15.pth --show-path=./exp  --include-values

7. 测试图片的命令

python tools/test.py configs/resnet/resnet50_8xb32_in1k.py  work_dirs/resnet50_8xb32_in1k/epoch_20.pth --work-dir=./exp --show-dir=./exp

测试结果会保存在exp文件夹中：

使用mmdeploy转换为engine模型。

mmdeploy官网：https://github.com/open-mmlab/mmdeploy

1.转换指令

python tools/deploy.py configs/mmpretrain/classification_onnxruntime_static.py ../mmpretrain/configs/resnet/resnet50_8xb32_in1k.py ../epoch_100.pth ../mmpretrain/data/gateDataset/test/close/151.jpg --work-dir=./work

这里的configs使用的是classification_onnxruntime_static.py，这个转换出来是onnx的形式。

python tools/deploy.py configs/mmpretrain/classification_tensorrt_static-224x224.py ../mmpretrain/configs/resnet/resnet50_8xb32_in1k.py ../epoch_100.pth ../mmpretrain/data/gateDataset/test/close/151.jpg --work-dir=./resnet50_gate --device cuda:0

这个转换出来会有engine形式，保存在resnet50_gate的文件夹下