图像修复模型MAT（Mask-Aware Transformer）的训练、推理实战记录

一、模型简介

        MAT（Mask-Aware Transformer for Large Hole Image Inpainting）是一个用于大孔洞图像修复的Transformer模型。该项目在CVPR 2022中获得了最佳论文提名，并且在图像修复领域取得了显著的成果。MAT模型通过结合Transformer和Mask机制，能够高效地修复图像中的大孔洞，生成高质量且多样化的修复结果。

        MAT项目的主要特点包括：

        （1）高保真度：生成的修复图像具有高度的真实感，减少了修复过程中产生的伪影。
        （2）多样性：模型能够生成多种可能的修复结果，增加了修复图像的多样性。
        （3）大孔洞修复：特别适用于修复图像中的大孔洞区域。

二、模型仓库文件

        项目解压缩之后的目录结构组织及文件如下：

MAT/
├── datasets               # 数据集处理相关脚本
│   └── ...
├── dnnlib                 # 深度学习网络库
│   └── ...
├── evaluatoin             # 评估脚本存放处
│   └── ...
├── figures                # 图表和可视化数据
│   └── ...
├── losses                 # 损失函数定义
│   └── ...
├── metrics                # 评价指标实现
│   └── ...
├── networks               # 网络架构定义
│   └── ...
├── test_sets              # 测试集图片和对应的mask
│   └── ...
├── torch_utils            # PyTorch相关的实用工具
│   └── ...
├── training               # 训练脚本及相关配置
│   └── ...
├── .gitignore             # Git忽略文件列表
├── LICENSE                # 项目许可协议
├── README.md              # 主要的项目说明文件
├── dataset_tool.py        # 数据集处理工具
├── generate_image.py      # 核心脚本，用于生成修复后的图像
├── legacy.py              # 兼容旧代码
├── requirements.txt       # 项目依赖清单
├── train.py               # 训练模型的主要脚本
└── ...

datasets：包含数据预处理和加载的脚本。dnnlib：深度神经网络的基础库。evaluatoin：用于评估模型性能的脚本。figures：存储实验结果的图表。losses：定义在训练过程中使用的损失函数。metrics：实现不同的评价指标，如FID，U-IDS等。networks：自定义的网络结构代码。test_sets：测试集的数据和对应的掩码。torch_utils：PyTorch特定的辅助功能。training：训练模型的相关配置和脚本。gitignore：Git忽略文件。LICENSE：许可证。README.md: 项目简介。dataset_tool.py：数据处理。generate_image.py：生成修复图像。train.py: 模型训练。

        虽然该项目没有单独列出的“配置文件”以传统意义上的.ini或.yaml形式存在，但配置主要是通过train.py中的命令行参数实现的。用户可以直接在调用train.py时指定各项训练细节，比如选用的网络结构(--generator, --discriminator)、损失函数(--loss)、优化策略等。这种灵活性使得配置通过参数传递的方式来进行，无需编辑独立的配置文件。

三、环境配置

        环境还有些问题

        按照官方的环境配置会有些问题，比如说pytorh出现警告等，后续看看能不能成功训练吧

四、模型训练

4.1 前言

         数据准备：确保输入图像和掩码图像的尺寸是512的倍数，如果不是，可以使用填充或调整大小的方式进行处理。但是我目前还不知道这个模型的数据集的组织结构是什么样的，所以训练的时候只能一种结构一种结构地进行尝试，还没有尝试成功，如果有人知道麻烦告知一下，不胜感激。
        模型选择：根据不同的数据集选择合适的预训练模型，例如CelebA-HQ、FFHQ或Places365-Standard。
        参数调整：根据具体需求调整生成图像的参数，如--truncation和--style_mix，以获得最佳的修复效果。

4.2 训练代码

        在Places数据集上训练模型，执行包含以下参数的bash脚本：

python train.py \--outdir=output_path \--gpus=8 \--batch=32 \--metrics=fid36k5_full \--data=training_data_path \--data_val=val_data_path \--dataloader=datasets.dataset_512.ImageFolderMaskDataset \--mirror=True \--cond=False \--cfg=places512 \--aug=noaug \--generator=networks.mat.Generator \--discriminator=networks.mat.Discriminator \--loss=losses.loss.TwoStageLoss \--pr=0.1 \--pl=False \--truncation=0.5 \--style_mix=0.5 \--ema=10 \--lr=0.001

参数解释：

outdir: 模型日志与权重的存储路径
gpus: 启用的GPU设备数量
batch: 多GPU集群的总批次大小（所有GPU处理的图像总数）
metrics: 完整评估指标集详见 metrics/metric_main.py （支持FID/KID等生成质量指标）
data: 训练集数据路径（需规范化为TFRecord格式）
data_val: 验证集数据路径（建议占总数据量的 10%-20%）
dataloader: 支持自定义数据加载模块（需继承BaseDataLoader类）
mirror: 镜像翻转数据增强开关（默认不启用）
cond: 类条件生成模式（需配合标签数据使用，默认关闭）
cfg: 模型架构配置基准（配置文件参考 train.py 中的BaseConfig类）
aug: 启用StyleGAN-ADA自适应数据增强（适用于小数据集场景）
generator: 自定义生成器网络（需兼容GANBaseGenerator接口）
discriminator: 自定义判别器网络（需实现梯度惩罚机制）
loss: 支持自定义损失函数（需返回梯度可微的标量值）
pr: 感知损失占比（需与VGG16特征提取器配合使用）
pl: 路径长度正则化策略（StyleGAN2技术，默认关闭）
truncation: 隐空间截断比例（取值范围 0.5-1.0，控制生成多样性）
style_mix: 隐变量风格混合概率（建议范围 0.6-0.9）
ema: 生成器权重指数滑动平均（建议设置K=10,000次迭代）
lr: Adam优化器的全局学习率（典型值 0.002-0.0002）

五、模型推理

        推理代码如下：

python generate_image.py --network model_path --dpath data_path --outdir out_path [--mpath mask_path]

        其中掩码路径为可选参数。若未指定，系统将自动生成512x512的随机掩码。需注意：掩码中的0值表示被遮蔽的像素区域，1值表示保留的像素区域。

python generate_image.py --network pretrained/CelebA-HQ.pkl --dpath test_sets/CelebA-HQ/images --mpath test_sets/CelebA-HQ/masks --outdir samples

        仅支持生成尺寸为512整数倍的图像。您需要通过填充(padding)或尺寸调整(resizing)使图像尺寸符合512整数倍要求。掩码填充请使用0值进行边缘扩展。

        若需使用CelebA-HQ-256模型，请在generate_image.py中显式设置'resolution'参数为256。