MMagic项目教程：如何设计自己的图像生成与编辑模型-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00920/article/details/148464645

MMagic项目教程：如何设计自己的图像生成与编辑模型

mmagic OpenMMLab Multimodal Advanced, Generative, and Intelligent Creation Toolbox. Unlock the magic 🪄: Generative-AI (AIGC), easy-to-use APIs, awsome model zoo, diffusion models, for text-to-image generation, image/video restoration/enhancement, etc. 项目地址: https://gitcode.com/gh_mirrors/mm/mmagic

概述

MMagic是一个基于MMEngine和MMCV构建的强大开源项目，专注于图像和视频的生成与编辑任务。本文将详细介绍如何在MMagic框架下设计自己的模型，包括超分辨率模型和生成对抗网络(GAN)的实现方法。

MMagic模型架构解析

在MMagic中，算法模型被清晰地分为两个核心组件：

Model（模型）：顶层封装，继承自MMEngine的BaseModel，负责完整的训练流程，包括前向传播、损失计算、反向传播和参数更新。
Module（模块）：构成模型的基础组件，包括：
- 网络架构（如生成器、判别器）
- 预定义的损失函数
- 数据预处理模块

这种分层设计使得模型开发更加模块化和可维护。

超分辨率模型SRCNN实现详解

1. 定义SRCNN网络结构

SRCNN是首个应用于单幅图像超分辨率的深度学习方法。我们通过继承BaseModule来实现：

@MODELS.register_module()
class MSRResNet(BaseModule):
    def __init__(self, in_channels, out_channels, mid_channels=64, 
                 num_blocks=16, upscale_factor=4):
        super().__init__()
        # 定义网络层
        self.conv_first = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
        self.trunk_net = make_layer(ResidualBlockNoBN, num_blocks, 
                                   mid_channels=mid_channels)
        # 上采样模块
        if upscale_factor == 4:
            self.upsample1 = PixelShufflePack(mid_channels, mid_channels, 2)
            self.upsample2 = PixelShufflePack(mid_channels, mid_channels, 2)
        # 输出层
        self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
        
    def forward(self, x):
        # 前向传播逻辑
        feat = self.conv_first(x)
        out = self.trunk_net(feat)
        out = self.upsample1(out)
        out = self.upsample2(out)
        return self.conv_last(out)

关键点：

使用@MODELS.register_module()装饰器注册网络
支持不同上采样比例因子
采用无BN的残差块结构

2. 构建完整的SRCNN模型

创建BaseEditModel作为顶层封装：

@MODELS.register_module()
class BaseEditModel(BaseModel):
    def __init__(self, generator, pixel_loss, train_cfg=None, test_cfg=None):
        super().__init__()
        self.generator = MODELS.build(generator)
        self.pixel_loss = MODELS.build(pixel_loss)
        
    def forward_train(self, batch_inputs, data_samples):
        pred = self.generator(batch_inputs)
        gt = torch.stack([d.gt_img.data for d in data_samples])
        loss = self.pixel_loss(pred, gt)
        return dict(loss=loss)

3. 配置与训练

创建配置文件srcnn_x4k915_g1_1000k_div2k.py：

model = dict(
    type='BaseEditModel',
    generator=dict(
        type='SRCNNNet',
        channels=(3, 64, 32, 3),
        kernel_sizes=(9, 1, 5),
        upscale_factor=4),
    pixel_loss=dict(type='L1Loss'))

启动训练：

python tools/train.py configs/srcnn/srcnn_x4k915_g1_1000k_div2k.py

DCGAN实现详解

1. 定义生成器和判别器

生成器实现：

@MODULES.register_module()
class DCGANGenerator(nn.Module):
    def __init__(self, output_scale=64, noise_size=100):
        super().__init__()
        # 从噪声到初始特征
        self.noise2feat = ConvModule(
            noise_size, 1024, 4, 1, 0, 
            conv_cfg=dict(type='ConvTranspose2d'))
        
        # 上采样模块
        self.upsampling = nn.Sequential(
            ConvModule(1024, 512, 4, 2, 1, 
                      conv_cfg=dict(type='ConvTranspose2d')),
            ConvModule(512, 256, 4, 2, 1,
                      conv_cfg=dict(type='ConvTranspose2d')))
        
        # 输出层
        self.output_layer = ConvModule(
            256, 3, 4, 2, 1,
            conv_cfg=dict(type='ConvTranspose2d'),
            act_cfg=dict(type='Tanh'))

判别器实现：

@MODULES.register_module()
class DCGANDiscriminator(nn.Module):
    def __init__(self, input_scale=64):
        super().__init__()
        # 下采样模块
        self.downsampling = nn.Sequential(
            ConvModule(3, 128, 4, 2, 1),
            ConvModule(128, 256, 4, 2, 1),
            ConvModule(256, 512, 4, 2, 1))
        
        # 输出层
        self.output_layer = ConvModule(
            512, 1, 4, 1, 0,
            act_cfg=None, norm_cfg=None)

2. 构建DCGAN模型

@MODELS.register_module()
class DCGAN(BaseModel):
    def __init__(self, generator, discriminator, gan_loss):
        super().__init__()
        self.generator = MODELS.build(generator)
        self.discriminator = MODELS.build(discriminator)
        self.gan_loss = MODELS.build(gan_loss)
        
    def forward_train(self, real_imgs, noise):
        # 生成假图像
        fake_imgs = self.generator(noise)
        
        # 判别器损失
        pred_real = self.discriminator(real_imgs)
        pred_fake = self.discriminator(fake_imgs.detach())
        loss_d = self.gan_loss(pred_real, pred_fake, True)
        
        # 生成器损失
        pred_fake = self.discriminator(fake_imgs)
        loss_g = self.gan_loss(pred_fake, None, False)
        
        return dict(loss_d=loss_d, loss_g=loss_g)

3. 训练配置

model = dict(
    type='DCGAN',
    generator=dict(
        type='DCGANGenerator',
        output_scale=64,
        noise_size=100),
    discriminator=dict(
        type='DCGANDiscriminator',
        input_scale=64),
    gan_loss=dict(type='GANLoss'))