MMSegmentation测试Segformer并统计模型参数量

#下载源码并解压缩

wget https://github.com/open-mmlab/mmsegmentation/archive/refs/tags/v1.2.1.tar.gz -O mmsegmentation-1.2.1.tar.gz
tar zxvf mmsegmentation-1.2.1.tar.gz

#安装mmcv

pip install mmcv==2.1.0

#安装openmim

pip3 install -U openmim

#安装其它依赖包

cd mmsegmentation-1.2.1
pip3 install -r requirements/runtime.txt
pip3 install -r requirements/tests.txt
pip3 install -r requirements/multimodal.txt
python setup.py develop -i https://pypi.tuna.tsinghua.edu.cn/simple

#下载配置文件和模型文件

mim download mmsegmentation --config segformer_mit-b2_8xb1-160k_cityscapes-1024x1024 --dest .

#测试SegFormer推理单张图片
 

python demo/image_demo.py demo/demo.png configs/segformer/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth --device cuda:0 --out-file result.jpg

#测试SegFormer推理多张图片

代码仿照demo/image_demo.py修改一下

import os
from glob import glob
from tqdm import tqdm
from argparse import ArgumentParser
from mmengine.model import revert_sync_batchnorm
from mmseg.apis import inference_model, init_model, show_result_pyplot


def main():
    parser = ArgumentParser()
    parser.add_argument('--input', help='Image file path')
    parser.add_argument('--config', help='Config file')
    parser.add_argument('--checkpoint', help='Checkpoint file')
    parser.add_argument('--output', help='Path to output file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--opacity',
        type=float,
        default=0.5,
        help='Opacity of painted segmentation map. In (0, 1] range.')
    parser.add_argument(
        '--title', default='result', help='The image identifier.')
    args = parser.parse_args()

    # build the model from a config file and a checkpoint file
    model = init_model(args.config, args.checkpoint, device=args.device)
    if args.device == 'cpu':
        model = revert_sync_batchnorm(model)

    image_files = []
    if os.path.exists(args.input):
        if os.path.isdir(args.input):
            for ext in ['png', 'jpg', 'jpeg', 'bmp']:
                files = glob(os.path.join(args.input, '**/*.%s' % (ext)), recursive=True)
                if len(files) > 0:
                    image_files.extend(files)
    else:
        return

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    for image_path in tqdm(image_files):
        (filename, extension) = os.path.splitext(image_path)
        filename = os.path.basename(filename)
        image_name = "result_" + filename
        output_path = os.path.join(args.output, image_name + ".jpg")

        # process a single image
        result = inference_model(model, image_path)
        # show the results
        show_result_pyplot(
            model,
            image_path,
            result,
            title=args.title,
            opacity=args.opacity,
            draw_gt=False,
            show=False,
            out_file=output_path)


if __name__ == '__main__':
    main()

#推理测试一个文件夹下的所有图片

python demo/images_dir_demo.py --input demo/images --config configs/segformer/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py --checkpoint segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth --device cuda:0 --output demo/result

统计模型参数量的代码如下：

from mmengine.analysis import get_model_complexity_info
from mmengine.analysis.print_helper import _format_size
from mmseg.apis import inference_model, init_model
from mmengine import Config
import torch


def load_model(config_path, checkpoint_path, resize_height, resize_width, class_num):
    cfg = Config.fromfile(config_path)
    cfg.crop_size = (resize_height, resize_width)
    cfg.data_preprocessor.size = cfg.crop_size
    cfg.model.data_preprocessor.size = cfg.crop_size
    cfg.test_pipeline = [
        dict(type='LoadImageFromFile'),
        dict(type='Resize', scale=(resize_width, resize_height), keep_ratio=True),
        # add loading annotation after ``Resize`` because ground truth
        # does not need to do resize data transform
        dict(type='LoadAnnotations'),
        dict(type='PackSegInputs')
    ]

    cfg.model.decode_head.num_classes = class_num
    print('class_num=', class_num)
    model = init_model(cfg, checkpoint_path, 'cuda:0')

    return model


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print('==> Building model..')
    config_path = 'segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py'
    checkpoint_path = 'segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth'
    # config_path = 'deeplabv3plus_r50b-d8_4xb2-80k_cityscapes-512x1024.py'
    # checkpoint_path = 'deeplabv3plus_r50b-d8_512x1024_80k_cityscapes_20201225_213645-a97e4e43.pth'
    resize_width = 1024
    resize_height = 512
    class_num = 19    # Cityscape是19个类别
    net = load_model(config_path, checkpoint_path, resize_height, resize_width, class_num)
    net = net.to(device=device)
    input_shape = (3, resize_height, resize_width)
    outputs = get_model_complexity_info(net, input_shape=input_shape, show_table=False, show_arch=False)
    flops = _format_size(outputs['flops'])
    params = _format_size(outputs['params'])
    print("flops:{}".format(flops))
    print("params:{}".format(params))