MMGeneration项目快速入门指南：模型推理与训练全流程解析

MMGeneration项目快速入门指南：模型推理与训练全流程解析

mmgeneration MMGeneration is a powerful toolkit for generative models, based on PyTorch and MMCV. 项目地址: https://gitcode.com/gh_mirrors/mm/mmgeneration

概述

MMGeneration是一个功能强大的生成对抗网络(GAN)框架，支持多种主流生成模型，包括无条件GAN、条件GAN、图像翻译模型等。本文将详细介绍如何使用MMGeneration进行模型推理、训练和评估，帮助开发者快速上手这一框架。

模型推理实践

无条件GAN图像生成

无条件GAN能够根据随机噪声生成逼真图像。MMGeneration提供了简洁的API实现这一功能：

from mmgen.apis import init_model, sample_unconditional_model

# 初始化模型
config = 'configs/styleganv2/stylegan2_c2_ffhq_1024_b4x8.py'
checkpoint = 'path_to_checkpoint.pth'
model = init_model(config, checkpoint, device='cuda:0')

# 生成4张图像
generated_images = sample_unconditional_model(model, 4)

也可以通过命令行工具实现：

python demo/unconditional_demo.py configs/styleganv2/stylegan2_c2_ffhq_1024_b4x8.py path_to_checkpoint.pth

条件GAN图像生成

条件GAN能够根据类别标签生成特定类别的图像：

from mmgen.apis import init_model, sample_conditional_model

model = init_model(config, checkpoint, device='cuda:0')

# 随机生成不同类别的图像
images = sample_conditional_model(model, 4)

# 生成特定类别的图像
cat_images = sample_conditional_model(model, 4, label=0)

命令行工具支持更灵活的生成方式：

python demo/conditional_demo.py config_file checkpoint_file \
    --label 0 1 2 3 \  # 指定生成类别
    --samples-per-classes 5  # 每类生成5张

图像翻译模型应用

图像翻译模型如Pix2Pix可以实现图像风格转换：

from mmgen.apis import sample_img2img_model

model = init_model(config, checkpoint)
translated_img = sample_img2img_model(model, 'input.jpg', target_domain='photo')

命令行使用方式：

python demo/translation_demo.py config_file checkpoint_file input.jpg

数据集准备规范

无条件GAN数据集

无条件GAN只需要包含真实图像的文件夹结构：

data/celeba/
├── img1.jpg
├── img2.jpg
└── ...

建议使用符号链接组织数据集：

mkdir data
ln -s /path/to/dataset ./data/dataset_name

图像翻译数据集

1. 配对数据集（如Pix2Pix）：

data/edges2shoes/
├── train
│   └── paired_img_AB.jpg  # 包含A和B域的拼接图像
└── test
    └── paired_img_AB.jpg

2. 非配对数据集（如CycleGAN）：

data/horse2zebra/
├── trainA
│   └── horse1.jpg
├── trainB
│   └── zebra1.jpg
├── testA
│   └── horse2.jpg
└── testB
    └── zebra2.jpg

模型训练指南

分布式训练

推荐使用分布式训练加速模型收敛：

# 常规环境
sh tools/dist_train.sh config_file 8  # 使用8块GPU

# Slurm集群环境
sh tools/slurm_train.sh partition job_name config_file work_dir

关键训练参数

• --work-dir: 指定工作目录保存日志和检查点
• --resume-from: 从检查点恢复训练
• --no-validate: 跳过验证阶段
• --seed: 设置随机种子保证可复现性

CPU训练（仅限调试）

export CUDA_VISIBLE_DEVICES=-1
python tools/train.py config_file

注意：动态GAN不支持CPU训练，且训练速度较慢，仅建议用于调试。

模型评估方法

MMGeneration支持6种评估指标：MS-SSIM、SWD、IS、FID、Precision&Recall和PPL。

评估配置示例

在配置文件中添加metrics部分：

metrics = dict(
    fid50k=dict(
        type='FID',
        num_images=50000,
        inception_pkl='path_to_stats.pkl',
        bgr2rgb=True))

执行评估

# 在线评估（不保存生成图像）
sh tools/dist_eval.sh config_file checkpoint_file 8 --online

# 离线评估（保存生成图像）
sh tools/dist_eval.sh config_file checkpoint_file 8 --samples-path outputs/

FID评估详解

FID是衡量生成图像质量的重要指标，MMGeneration提供两种实现：

1. PyTorch版本：使用修改的InceptionV3网络
2. Tero版本：使用TensorFlow InceptionV3（需PyTorch≥1.6）

提取真实图像特征的预处理：

python tools/utils/inception_stat.py \
    --imgsdir real_images/ \
    --pklname real_stats.pkl \
    --size 256  # 图像尺寸

实用技巧

1. 多GPU评估可显著加速FID和IS计算
2. 使用--eval none可仅生成图像不计算指标
3. 图像翻译评估需指定目标域：

   python tools/utils/translation_eval.py config_file checkpoint_file --t photo
   

通过本指南，开发者可以快速掌握MMGeneration的核心功能，包括模型推理、训练和评估的全流程。该框架的模块化设计使得各种生成模型的实现和评估变得简单高效。

mmgeneration MMGeneration is a powerful toolkit for generative models, based on PyTorch and MMCV. 项目地址: https://gitcode.com/gh_mirrors/mm/mmgeneration