简介图像修复是计算机视觉中极具实用价值的研究方向旨在通过算法自动恢复图像中缺失、遮挡或破损区域的像素内容。传统方法基于纹理合成或插值在面对大面积缺失时往往力不从心而基于深度学习的生成模型则能借助语义理解推断出合理且自然的内容。以生成对抗网络GAN为核心、U-Net为生成器主干配合感知损失与对抗损失的联合优化已成为当前主流技术范式。这类系统可用于老照片修复、物体移除、影视后期及医学影像处理等真实场景。在实际工程落地中基于PyTorch搭建的修复项目需要重点关注数据Mask生成、生成器与判别器结构、损失函数配比以及训练推理流程的稳定性。本文围绕一套完整的基于PyTorch的图像修复源码系统梳理其设计思路、环境配置、核心模块实现与训练推理细节并针对常见问题给出排查方法适合需要学习或二次开发相关系统的开发者参考。 前两天清理硬盘翻出一个标注为“基于PyTorch的图像修复系统”的源码压缩包顺手解压跑了一轮。图像修复Image Inpainting是计算机视觉里一个很实用的方向输入一张被遮挡、划痕或物体缺失的图片模型会自动把缺失区域补出来比传统克隆印章和插值算法自然得多。这类源码适合正在学PyTorch的开发者研究也适合有图像修复需求的产品工程师做二次开发。这篇文章就把整个系统的设计思路、环境搭建、核心代码和复现过程中的坑完整梳理一遍。这类项目在网上有不少公开实现但大多只给了模型结构训练流程写得模糊数据预处理也不完整。我拿到这个压缩包后重点检查了三个东西数据加载和Mask生成逻辑、生成器和判别器的实现细节、训练与推理的入口配置。只要这三个地方能跑通整个系统基本就能复现。下面从头开始拆。1. 项目整体设计与核心思路1.1 图像修复任务与适用场景图像修复的核心任务是给定一张带有缺失区域的图像以及标明缺失位置的Mask模型需要生成与周围上下文一致的像素内容。用数学语言描述输入是原始图像 (I) 和掩码 (M)其中 (M1) 的位置代表缺失区域修复目标是估计 (I_{out})使得 (I_{out}) 在 Mask 区域内与真实内容 (I_{gt}) 尽可能接近同时在视觉上没有明显接缝。传统方法比如Telea算法和基于Patch匹配的纹理合成对细小划痕处理还可以遇到大块缺失区域就无能为力要么模糊要么纹理重复。深度学习模型能够依据高层语义信息推断出合理内容比如补全一栋被遮挡的楼、一条被电线穿过的天空甚至生成原来并不确定的细节。实际落地场景主要有几类老照片修复去除照片上的折痕、污渍、霉斑同时恢复背景纹理。图像编辑把不需要的元素路人、水印、杂物从画面中移除再用修复算法填充背景。影视后期对拍摄时无法规避的穿帮物体做内容补全。医学影像去除扫描图像中的金属伪影或运动伪影。这个源码的通用性还不错只要你准备了合适的Mask分布和数据集微调一下就能适配上述场景。1.2 源码模块结构与架构选型解压压缩包后典型工程结构大致如下. ├── checkpoints/ # 模型权重保存位置 ├── configs/ # 训练和推理参数配置文件 ├── data/ # 数据集加载、Mask生成、数据增强 ├── models/ # 生成器、判别器、损失函数定义 ├── utils/ # 图像处理工具、评估指标 ├── train.py # 训练入口 ├── infer.py # 推理入口 ├── requirements.txt # 依赖列表 └── README.md # 使用说明拿到源码千万别急着训练先看README和requirements确认作者用了哪个PyTorch版本、数据是什么格式、训练输入尺寸是多大。很多人复现失败问题不在模型而在版本和配置不一致。比如PyTorch 1.x和2.x在某些算子行为上差异不小模型代码里如果用了torch.nn.functional.interpolate的align_corners参数版本不同结果可能完全不同。架构选型上这个项目遵循了主流修复框架的设计生成器用U-Net变体配合部分卷积或门控卷积判别器用PatchGAN损失函数由像素重建损失、感知损失和对抗损失组成。这样组合的原因是重建损失让网络学到稳定的基础结构感知损失从特征层面保证语义一致对抗损失负责让修复区域纹理更真实。2. 环境准备与依赖安装2.1 PyTorch环境搭建的完整步骤先强调一件事图像修复训练必须GPU纯CPU跑会慢得让人怀疑人生。配置环境我建议用Anaconda管理虚拟环境不要直接装在base环境里否则后面项目一多依赖冲突能烦死你。创建独立环境并激活conda create -n inpaint python3.8 -y conda activate inpaint接下来安装GPU版PyTorch。先看自己机器的CUDA支持情况nvidia-smi输出右上角能看到驱动支持的最高CUDA版本比如12.1那么安装PyTorch时选择cuda 12.1或比它低的版本都可以。注意nvidia-smi显示的CUDA版本是驱动支持的版本不是当前环境已经装好的版本。实际安装PyTorch时conda会自动把配套的CUDA runtime和cuDNN一起装进虚拟环境所以不需要单独装CUDA Toolkit除非你要编译自定义算子。安装命令示例conda install pytorch torchvision pytorch-cuda11.8 -c pytorch -c nvidia如果下载速度不稳定可以用清华源加速但要注意conda和pip的源不要混着换出问题。我这里不展开说镜像配置重点是你的pytorch-cuda版本要和显卡驱动兼容。装完之后验证python -c import torch; print(torch.__version__, torch.cuda.is_available())输出True说明CUDA可用否则要重新检查安装步骤。2.2 依赖库安装与预训练权重准备requirements.txt里面一般会有这些库torch torchvision opencv-python numpy pillow tqdm tensorboard scikit-image批量安装pip install -r requirements.txt如果需求里有pytorch-msssim、lpips这类评估指标库建议一并装上后面测试效果会用到。安装时如果遇到opencv-python编译慢可以直接用阿里云或豆瓣的镜像源安装。数据集方面如果只是想快速跑通流程不需要一上来就下载Places2这么大的数据集。可以先拿CelebA或COCO的子集试跑甚至用自己拍摄的几十张照片也能验证。但需要注意修复模型对数据量有要求数据太少容易过拟合表现就是训练loss很低换一张新图效果崩。公开数据集我常用Places2和CelebA-HQ前者适合场景补全后者适合人脸修复。预训练权重一般放在checkpoints目录。加载权重时报size mismatch是最常见的问题原因通常是模型结构和state_dict键名对不上。遇到这种情况先用torch.load把权重load进来打印model_state_dict的keys和当前模型的keys做对比缺哪个补哪个多了的删掉再加载就顺畅了。2.3 训练与推理的关键参数解析打开配置文件常见参数如下表所示参数推荐值说明image_size256 / 512输入图像尺寸越大越消耗显存batch_size4 ~ 8显存不足时优先调低learning_rate1e-4 ~ 2e-4Adam优化器常用范围lambda_rec10重建损失权重lambda_perceptual0.1感知损失权重lambda_adv1对抗损失权重iterations100000按迭代数训练更常见save_interval5000每隔多少步保存一次模型这里重点说下损失权重的影响。重建损失权重太高模型倾向于输出平滑结果细节会糊对抗损失权重太高训练不稳定容易出现色彩失真。我用过的组合里先固定lambda_rec10, lambda_perceptual0.1再逐步从0.1调到1效果会更可控。训练过程中需要同时关注生成器和判别器的loss比例判别器loss长期接近0说明生成器完全打不过判别器梯度几乎没有需要降低判别器学习率。3. 核心模块源码解析与实操3.1 数据加载与Mask生成逻辑数据加载是第一个容易踩坑的地方。Pytorch的Dataset类返回的样本格式必须和模型输入对齐。修复任务中输入不是单一图像而是“损坏图像 Mask”的组合。很多实现会在__getitem__里做以下操作读取完整图像并做随机裁剪或resize。生成随机Mask。根据Mask对图像做损坏处理比如把Mask区域像素置为0。将Mask归一化为0/1并与图像在通道维度拼接得到4通道输入。返回(input_tensor, mask_tensor, gt_tensor, mask_for_loss)。Mask生成是重点。如果Mask全是随机矩形模型只会补矩形区域遇到真实划痕就失效。我在项目里看到比较实用的做法是生成三种类型随机矩形块模拟物体遮挡。随机线条和曲线模拟划痕。不规则多边形模拟污渍。实现时可以用OpenCV画线、画多边形再配合膨胀腐蚀让Mask边缘更自然。简单示例import cv2 import numpy as np def random_mask(height, width): mask np.zeros((height, width), dtypenp.uint8) # 随机矩形 x, y, w, h np.random.randint(0, width//3, 4) mask[y:yh, x:xw] 255 # 随机线条 pts np.random.randint(0, height, (2, 2)) cv2.line(mask, tuple(pts[0]), tuple(pts[1]), 255, 10) mask cv2.dilate(mask, np.ones((5, 5), np.uint8)) return mask / 255.0这只是示例正式项目里会加入更多形态变化。Mask生成时一定要保证训练和推理的Mask分布一致。如果训练时Mask区域都是小面积推理时给一个大面积Mask模型就会表现得很差。3.2 生成器与判别器的网络实现细节生成器最基础的做法是把U-Net输入改成4通道中间换几个残差块输出3通道。但标准卷积在处理Mask区域时有天然缺陷卷积核对所有像素一视同仁缺失区域的零像素会污染特征导致修复结果有灰斑和边界模糊。所以成熟的修复项目会用部分卷积或门控卷积。部分卷积的核心思路是在卷积操作时只对有效像素做计算让Mask区域不参与特征更新。输出是这样算的out W * (X * M) / sum(M) b mask_out 1 if sum(M) 0 else 0每一步卷积后Mask也要跟着更新这样网络能自动判断哪些位置已经被修复哪些还需要继续生成。用PyTorch自定义部分卷积层要注意实现细节比如分母sum(M)不能为0需要加一个epsilon。判别器用PatchGAN。简单来说它不是输出一个全局真/假标量而是输出一个特征图比如16x16的矩阵每个值代表输入图像局部区域是真还是假。这么做可以让判别器关注局部纹理一致性避免出现“整体像局部崩”的情况。PatchGAN实现并不复杂import torch.nn as nn class PatchDiscriminator(nn.Module): def __init__(self, in_channels3): super().__init__() self.layers nn.Sequential( nn.Conv2d(in_channels, 64, 4, 2, 1), nn.LeakyReLU(0.2), nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2), nn.Conv2d(128, 256, 4, 2, 1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2), nn.Conv2d(256, 1, 4, 1, 1) ) def forward(self, x): return self.layers(x)实际项目里输入可能是4通道含Mask需要微调第一层输入尺寸。3.3 损失函数组合与训练循环模型的优化目标不是一个loss而是多个loss的加权和。典型组合如下loss_rec l1_loss(pred, gt, mask) # 只计算Mask内区域 loss_perceptual perceptual_loss(pred, gt) # VGG特征距离 loss_adv gan_loss(discriminator(pred), real_label) total_loss loss_rec * w_rec loss_perceptual * w_perceptual loss_adv * w_adv重建损失用L1比L2好L2会过度惩罚大误差导致输出偏向平均颜色边缘会变模糊。L1对离群值更鲁棒能保留更多细节。感知损失一般用ImageNet预训练的VGG16取relu1_1, relu2_1, relu3_1, relu4_1几层特征计算生成图与真实图特征之间的L1距离。注意使用预训练VGG时输入图像要做相同的归一化否则特征值分布不同感知损失的意义会打折扣。训练循环里生成器和判别器交替更新。一个常见的策略是每个iteration里先更新生成器再更新判别器或者每更新一次生成器更新两次判别器。这个项目源码里如果没有控制判别器更新频率训练出现震荡可以自己加一个if step % 2 0的更新判别器逻辑。训练日志建议记录step, total_loss, rec_loss, perc_loss, adv_loss, d_loss, psnr, ssim伪代码可以这样写for step in range(total_iterations): real_img, real_gt, mask next(data_loader) masked_img real_img * (1 - mask) # 训练生成器 fake_img generator(torch.cat([masked_img, mask], dim1)) rec_loss l1_loss(fake_img, real_gt, mask) perc_loss vgg_loss(fake_img, real_gt) adv_loss gan_loss(discriminator(fake_img), real_label) g_loss rec_loss * w_rec perc_loss * w_perc adv_loss * w_adv g_optimizer.zero_grad() g_loss.backward() g_optimizer.step() # 训练判别器 real_pred discriminator(real_gt) fake_pred discriminator(fake_img.detach()) d_loss (gan_loss(real_pred, real_label) gan_loss(fake_pred, fake_label)) * 0.5 d_optimizer.zero_grad() d_loss.backward() d_optimizer.step()这里面有一个很容易忽略的细节计算重建损失时建议把mask也作为权重传入只计算Mask区域内的像素差。如果把全图都算进去背景区域占了绝对主导模型很容易学到“把背景复制一下就行”对被遮挡区域毫无生成能力。4. 实战从零训练与图像修复推理4.1 用自己的数据集跑通训练建议第一次跑用一个小数据集比如从公开数据集里挑500张图片设置训练步数1000步目标不是效果好而是验证整个链路是通的数据加载正确、模型前向正常、损失能下降、权重能保存和加载。这一步通过后再全量训练能省下大量排查时间。使用自己的数据集时先创建一个data_list.txt每一行是图片路径/path/to/train/000001.jpg /path/to/train/000002.jpg ...然后修改config把data_root指向这个txt设置好image_size、batch_size等参数。执行训练python train.py --config configs/train_config.yaml训练启动后观察前几个迭代的日志。正常情况下loss应该在稳步下降PSNR和SSIM缓慢上升。如果loss曲线剧烈震荡我的经验是先把学习率调低一个数量级比如从1e-4调到1e-5再看是否稳定。如果生成器loss和判别器loss走势极不平衡参考上一节提到的调整训练频率。训练时长方面单张RTX 3090跑256×256输入、batch_size810万步大约需要2天。如果想快速验证可以先用--max_iters 2000跑一小段。4.2 推理流程与效果对比推理流程比训练简单但精度问题同样不可忽视。命令类似python infer.py --image samples/test.jpg --mask samples/mask.png --checkpoint output/latest.pth --output result.png推理脚本内部一般做这四步读取图像和Mask统一resize到训练尺寸。图像和Mask拼接经过生成器前向计算。生成结果与原始图像融合保留原始图中未损坏区域。保存输出图像必要时做后处理。这里的关键是融合。不能把生成器的整个输出直接作为结果因为它在非Mask区域也会产生偏移导致原图背景被改。正确的融合方式是result original_image * (1 - mask) generated_image * mask如果发现修复区域边缘有接缝可以对Mask做一次高斯模糊或者把Mask膨胀几个像素让融合过渡更自然。颜色偏色问题时检查推理时是否做了和训练一样的归一化。很多项目训练时把像素映射到[-1,1]推理时忘了减均值除方差出来的图就会明显偏色。效果对比时我习惯把原图、Mask、修复结果拼在一张画布里方便肉眼评估。重点关注三块边缘过渡是否自然、纹理是否重复、颜色是否一致。4.3 评估指标PSNR/SSIM如何看在有真实完整图像的测试集上可以计算PSNR和SSIM。PSNR越高越好通常修复模型能到25~35dBSSIM越接近1越好。另一个更接近主观感受的指标是LPIPS值越低越好。但客观指标和人的感受经常不一致。GAN生成的结果纹理细腻但像素级偏移大PSNR可能反而不如模糊的结果。因此评估时必须同时看主观效果图不能只看分数。我自己的习惯是先用PSNR/SSIM筛掉明显不行的模型再用肉眼对比候选模型的修复图最后根据应用场景做决定。如果业务场景对真实性要求高LPIPS和人工评测优先级要提前。5. 常见问题与排查技巧实录5.1 训练不收敛或损失异常训练时遇到loss变成NaN最先检查三件事输入图像有没有全黑或全白、Mask是否全0、学习率是否过大。图像修复模型输入是全零区域时前向计算容易出现梯度异常。可以先把学习率降到1e-5跑一个iteration看看如果还NaN就要检查数据归一化和模型权重初始化。另一个常见情况是loss快速下降但不代表效果好。生成器可能学会了一种投机取巧的方式非Mask区域直接复制Mask区域输出平均色。这样重建loss很低PSNR也可能不差但视觉上一塌糊涂。解决办法是提高感知损失和对抗损失的权重并定期人工查看生成图。我遇到的训练崩溃还有一个隐藏原因用FP16混合精度训练时loss数值不稳定。如果源码默认开了amp可以关闭再试或者调整grad_scaler的初始化比例。修复模型对精度比较敏感稳定优先于速度。5.2 修复结果模糊或伪影严重模糊是因为模型没有高频细节常见原因包括重建损失权重过高、网络容量不足、输入分辨率过低。我的做法是先降低lambda_rec同时确保lambda_perceptual不是0否则模型会大量丢失纹理。如果还模糊考虑换更大的生成器或增加中间通道数。伪影一般出现在GAN训练不稳定时。表现为修复区域有斑点、彩色条纹或过锐利的边缘。可以考虑降低对抗损失权重。使用LSGAN或Hinge Loss代替标准BCE。判别器增加谱归一化。训练前期冻结判别器只训练生成器若干千步。伪影还可能与推理分辨率有关。如果训练是256×256推理时给512×512的图模型没见过这么大尺寸容易产生结构畸变。此时优先保证推理尺寸和训练一致再考虑是否使用支持任意分辨率的结构。5.3 环境兼容性与显存不足问题环境兼容性问题中最常见的是PyTorch和CUDA版本不匹配。安装PyTorch时如果选了比驱动支持更高的CUDA版本会直接报CUDA driver version is insufficient。解决办法是用nvidia-smi确认支持版本然后重新安装匹配的PyTorch。显存不足的报错形式一般是RuntimeError: CUDA out of memory. Tried to allocate X MiB我推荐的排查顺序降低batch_size到2甚至1。降低图像尺寸比如从512降到256。设置torch.backends.cudnn.benchmarkTrue有时候能省显存。使用梯度累积模拟大batch。开启torch.utils.checkpoint梯度检查点用计算换显存。使用混合精度训练FP16。注意DataLoader的num_workers调大并不会减少显存占用反而可能因为多进程缓存导致整体内存上涨。显存不足时可以先关掉验证集评估因为验证过程也占用显存。5.4 常见问题速查表问题可能原因解决办法loss为NaN学习率过大或输入包含无效值降低学习率检查数据归一化修复区域模糊重建损失权重过高或网络容量不足降低lambda_rec增加网络宽度边界有接缝推理融合未处理Mask对Mask做膨胀或高斯模糊后融合颜色偏色推理归一化与训练不一致统一预处理逻辑CUDA out of memorybatch_size/分辨率过大调小尺寸使用FP16梯度累积预训练权重加载失败模型结构和权重键名不匹配打印keys逐一对比手动修改state_dict大型Mask修复效果差训练时Mask面积分布不匹配增加不规则大Mask的数据增强6. 项目扩展方向与个人心得6.1 如何扩展到自己的业务场景单纯跑通一个开源项目不算本事能把它用起来才是目标。假设你要做老照片修复需要自己准备一批带划痕的图片或者模拟划痕生成训练数据要移除水印就得专门生成文字区域的Mask。千万不要通用模型一把梭效果大概率不如预期。扩展方面可以考虑模型导出。PyTorch模型可以用ONNX导出再用TensorRT做推理加速方便部署到服务端torch.onnx.export( generator, (dummy_img, dummy_mask), inpaint.onnx, input_names[image, mask], output_names[output], dynamic_axes{image: {0: batch}, mask: {0: batch}} )ONNX导出时要注意自定义层比如部分卷积中的Mask更新可能不支持导出需要把部分卷积改写成兼容算子或者用torch.onnx.export时设置custom_opsets。这块我踩过坑建议先导出一个最简模型验证通道。另外把修复系统包装成API接口也很实用。用FastAPI加载模型接收图像和Mask返回修复结果这样就能接入小程序、网页或后端服务。要注意的是并发场景下显存可能不够可以考虑单进程单GPU队列方式避免多个请求同时占用显存导致OOM。6.2 实操后的几点建议最后说几点我反复踩坑之后总结出来的经验。第一个是“先复现再创新”。拿到源码后不要着急改结构先用作者给的预训练权重和测试命令跑通确认整个链路能出成果再开始替换模型、改损失函数否则出现问题根本分不清是原代码的锅还是你自己改的锅。第二个是“每次只改一个变量”。图像修复模型效果受多个因素影响如果一次改了三四个参数效果变差了你完全不知道是哪个参数导致的。我习惯用实验管理工具记录每次训练的配置、log和输出图比如简单地把每个实验放在单独目录命名带上参数摘要。第三个是保存模型时把配置一起保存。只存state_dict有个问题过几个月你自己都忘了当时用了什么image_size、什么损失权重、什么归一化方式。最好是存一个字典包含model_state_dict、config和源码版本方便日后加载和新环境复现。在实际修复效果上别迷信单模型很多时候“检测Mask 修复 后处理”的pipeline比单独调模型更有效。先用分割模型自动生成Mask修复后再用锐化或颜色校正提升观感。这个思路能大幅提高项目的落地价值也是我把这个源码项目吃透之后最想分享的一点。本文还有配套的精品资源点击获取
基于PyTorch的图像修复系统：从原理到源码复现实战解析
简介图像修复是计算机视觉中极具实用价值的研究方向旨在通过算法自动恢复图像中缺失、遮挡或破损区域的像素内容。传统方法基于纹理合成或插值在面对大面积缺失时往往力不从心而基于深度学习的生成模型则能借助语义理解推断出合理且自然的内容。以生成对抗网络GAN为核心、U-Net为生成器主干配合感知损失与对抗损失的联合优化已成为当前主流技术范式。这类系统可用于老照片修复、物体移除、影视后期及医学影像处理等真实场景。在实际工程落地中基于PyTorch搭建的修复项目需要重点关注数据Mask生成、生成器与判别器结构、损失函数配比以及训练推理流程的稳定性。本文围绕一套完整的基于PyTorch的图像修复源码系统梳理其设计思路、环境配置、核心模块实现与训练推理细节并针对常见问题给出排查方法适合需要学习或二次开发相关系统的开发者参考。 前两天清理硬盘翻出一个标注为“基于PyTorch的图像修复系统”的源码压缩包顺手解压跑了一轮。图像修复Image Inpainting是计算机视觉里一个很实用的方向输入一张被遮挡、划痕或物体缺失的图片模型会自动把缺失区域补出来比传统克隆印章和插值算法自然得多。这类源码适合正在学PyTorch的开发者研究也适合有图像修复需求的产品工程师做二次开发。这篇文章就把整个系统的设计思路、环境搭建、核心代码和复现过程中的坑完整梳理一遍。这类项目在网上有不少公开实现但大多只给了模型结构训练流程写得模糊数据预处理也不完整。我拿到这个压缩包后重点检查了三个东西数据加载和Mask生成逻辑、生成器和判别器的实现细节、训练与推理的入口配置。只要这三个地方能跑通整个系统基本就能复现。下面从头开始拆。1. 项目整体设计与核心思路1.1 图像修复任务与适用场景图像修复的核心任务是给定一张带有缺失区域的图像以及标明缺失位置的Mask模型需要生成与周围上下文一致的像素内容。用数学语言描述输入是原始图像 (I) 和掩码 (M)其中 (M1) 的位置代表缺失区域修复目标是估计 (I_{out})使得 (I_{out}) 在 Mask 区域内与真实内容 (I_{gt}) 尽可能接近同时在视觉上没有明显接缝。传统方法比如Telea算法和基于Patch匹配的纹理合成对细小划痕处理还可以遇到大块缺失区域就无能为力要么模糊要么纹理重复。深度学习模型能够依据高层语义信息推断出合理内容比如补全一栋被遮挡的楼、一条被电线穿过的天空甚至生成原来并不确定的细节。实际落地场景主要有几类老照片修复去除照片上的折痕、污渍、霉斑同时恢复背景纹理。图像编辑把不需要的元素路人、水印、杂物从画面中移除再用修复算法填充背景。影视后期对拍摄时无法规避的穿帮物体做内容补全。医学影像去除扫描图像中的金属伪影或运动伪影。这个源码的通用性还不错只要你准备了合适的Mask分布和数据集微调一下就能适配上述场景。1.2 源码模块结构与架构选型解压压缩包后典型工程结构大致如下. ├── checkpoints/ # 模型权重保存位置 ├── configs/ # 训练和推理参数配置文件 ├── data/ # 数据集加载、Mask生成、数据增强 ├── models/ # 生成器、判别器、损失函数定义 ├── utils/ # 图像处理工具、评估指标 ├── train.py # 训练入口 ├── infer.py # 推理入口 ├── requirements.txt # 依赖列表 └── README.md # 使用说明拿到源码千万别急着训练先看README和requirements确认作者用了哪个PyTorch版本、数据是什么格式、训练输入尺寸是多大。很多人复现失败问题不在模型而在版本和配置不一致。比如PyTorch 1.x和2.x在某些算子行为上差异不小模型代码里如果用了torch.nn.functional.interpolate的align_corners参数版本不同结果可能完全不同。架构选型上这个项目遵循了主流修复框架的设计生成器用U-Net变体配合部分卷积或门控卷积判别器用PatchGAN损失函数由像素重建损失、感知损失和对抗损失组成。这样组合的原因是重建损失让网络学到稳定的基础结构感知损失从特征层面保证语义一致对抗损失负责让修复区域纹理更真实。2. 环境准备与依赖安装2.1 PyTorch环境搭建的完整步骤先强调一件事图像修复训练必须GPU纯CPU跑会慢得让人怀疑人生。配置环境我建议用Anaconda管理虚拟环境不要直接装在base环境里否则后面项目一多依赖冲突能烦死你。创建独立环境并激活conda create -n inpaint python3.8 -y conda activate inpaint接下来安装GPU版PyTorch。先看自己机器的CUDA支持情况nvidia-smi输出右上角能看到驱动支持的最高CUDA版本比如12.1那么安装PyTorch时选择cuda 12.1或比它低的版本都可以。注意nvidia-smi显示的CUDA版本是驱动支持的版本不是当前环境已经装好的版本。实际安装PyTorch时conda会自动把配套的CUDA runtime和cuDNN一起装进虚拟环境所以不需要单独装CUDA Toolkit除非你要编译自定义算子。安装命令示例conda install pytorch torchvision pytorch-cuda11.8 -c pytorch -c nvidia如果下载速度不稳定可以用清华源加速但要注意conda和pip的源不要混着换出问题。我这里不展开说镜像配置重点是你的pytorch-cuda版本要和显卡驱动兼容。装完之后验证python -c import torch; print(torch.__version__, torch.cuda.is_available())输出True说明CUDA可用否则要重新检查安装步骤。2.2 依赖库安装与预训练权重准备requirements.txt里面一般会有这些库torch torchvision opencv-python numpy pillow tqdm tensorboard scikit-image批量安装pip install -r requirements.txt如果需求里有pytorch-msssim、lpips这类评估指标库建议一并装上后面测试效果会用到。安装时如果遇到opencv-python编译慢可以直接用阿里云或豆瓣的镜像源安装。数据集方面如果只是想快速跑通流程不需要一上来就下载Places2这么大的数据集。可以先拿CelebA或COCO的子集试跑甚至用自己拍摄的几十张照片也能验证。但需要注意修复模型对数据量有要求数据太少容易过拟合表现就是训练loss很低换一张新图效果崩。公开数据集我常用Places2和CelebA-HQ前者适合场景补全后者适合人脸修复。预训练权重一般放在checkpoints目录。加载权重时报size mismatch是最常见的问题原因通常是模型结构和state_dict键名对不上。遇到这种情况先用torch.load把权重load进来打印model_state_dict的keys和当前模型的keys做对比缺哪个补哪个多了的删掉再加载就顺畅了。2.3 训练与推理的关键参数解析打开配置文件常见参数如下表所示参数推荐值说明image_size256 / 512输入图像尺寸越大越消耗显存batch_size4 ~ 8显存不足时优先调低learning_rate1e-4 ~ 2e-4Adam优化器常用范围lambda_rec10重建损失权重lambda_perceptual0.1感知损失权重lambda_adv1对抗损失权重iterations100000按迭代数训练更常见save_interval5000每隔多少步保存一次模型这里重点说下损失权重的影响。重建损失权重太高模型倾向于输出平滑结果细节会糊对抗损失权重太高训练不稳定容易出现色彩失真。我用过的组合里先固定lambda_rec10, lambda_perceptual0.1再逐步从0.1调到1效果会更可控。训练过程中需要同时关注生成器和判别器的loss比例判别器loss长期接近0说明生成器完全打不过判别器梯度几乎没有需要降低判别器学习率。3. 核心模块源码解析与实操3.1 数据加载与Mask生成逻辑数据加载是第一个容易踩坑的地方。Pytorch的Dataset类返回的样本格式必须和模型输入对齐。修复任务中输入不是单一图像而是“损坏图像 Mask”的组合。很多实现会在__getitem__里做以下操作读取完整图像并做随机裁剪或resize。生成随机Mask。根据Mask对图像做损坏处理比如把Mask区域像素置为0。将Mask归一化为0/1并与图像在通道维度拼接得到4通道输入。返回(input_tensor, mask_tensor, gt_tensor, mask_for_loss)。Mask生成是重点。如果Mask全是随机矩形模型只会补矩形区域遇到真实划痕就失效。我在项目里看到比较实用的做法是生成三种类型随机矩形块模拟物体遮挡。随机线条和曲线模拟划痕。不规则多边形模拟污渍。实现时可以用OpenCV画线、画多边形再配合膨胀腐蚀让Mask边缘更自然。简单示例import cv2 import numpy as np def random_mask(height, width): mask np.zeros((height, width), dtypenp.uint8) # 随机矩形 x, y, w, h np.random.randint(0, width//3, 4) mask[y:yh, x:xw] 255 # 随机线条 pts np.random.randint(0, height, (2, 2)) cv2.line(mask, tuple(pts[0]), tuple(pts[1]), 255, 10) mask cv2.dilate(mask, np.ones((5, 5), np.uint8)) return mask / 255.0这只是示例正式项目里会加入更多形态变化。Mask生成时一定要保证训练和推理的Mask分布一致。如果训练时Mask区域都是小面积推理时给一个大面积Mask模型就会表现得很差。3.2 生成器与判别器的网络实现细节生成器最基础的做法是把U-Net输入改成4通道中间换几个残差块输出3通道。但标准卷积在处理Mask区域时有天然缺陷卷积核对所有像素一视同仁缺失区域的零像素会污染特征导致修复结果有灰斑和边界模糊。所以成熟的修复项目会用部分卷积或门控卷积。部分卷积的核心思路是在卷积操作时只对有效像素做计算让Mask区域不参与特征更新。输出是这样算的out W * (X * M) / sum(M) b mask_out 1 if sum(M) 0 else 0每一步卷积后Mask也要跟着更新这样网络能自动判断哪些位置已经被修复哪些还需要继续生成。用PyTorch自定义部分卷积层要注意实现细节比如分母sum(M)不能为0需要加一个epsilon。判别器用PatchGAN。简单来说它不是输出一个全局真/假标量而是输出一个特征图比如16x16的矩阵每个值代表输入图像局部区域是真还是假。这么做可以让判别器关注局部纹理一致性避免出现“整体像局部崩”的情况。PatchGAN实现并不复杂import torch.nn as nn class PatchDiscriminator(nn.Module): def __init__(self, in_channels3): super().__init__() self.layers nn.Sequential( nn.Conv2d(in_channels, 64, 4, 2, 1), nn.LeakyReLU(0.2), nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2), nn.Conv2d(128, 256, 4, 2, 1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2), nn.Conv2d(256, 1, 4, 1, 1) ) def forward(self, x): return self.layers(x)实际项目里输入可能是4通道含Mask需要微调第一层输入尺寸。3.3 损失函数组合与训练循环模型的优化目标不是一个loss而是多个loss的加权和。典型组合如下loss_rec l1_loss(pred, gt, mask) # 只计算Mask内区域 loss_perceptual perceptual_loss(pred, gt) # VGG特征距离 loss_adv gan_loss(discriminator(pred), real_label) total_loss loss_rec * w_rec loss_perceptual * w_perceptual loss_adv * w_adv重建损失用L1比L2好L2会过度惩罚大误差导致输出偏向平均颜色边缘会变模糊。L1对离群值更鲁棒能保留更多细节。感知损失一般用ImageNet预训练的VGG16取relu1_1, relu2_1, relu3_1, relu4_1几层特征计算生成图与真实图特征之间的L1距离。注意使用预训练VGG时输入图像要做相同的归一化否则特征值分布不同感知损失的意义会打折扣。训练循环里生成器和判别器交替更新。一个常见的策略是每个iteration里先更新生成器再更新判别器或者每更新一次生成器更新两次判别器。这个项目源码里如果没有控制判别器更新频率训练出现震荡可以自己加一个if step % 2 0的更新判别器逻辑。训练日志建议记录step, total_loss, rec_loss, perc_loss, adv_loss, d_loss, psnr, ssim伪代码可以这样写for step in range(total_iterations): real_img, real_gt, mask next(data_loader) masked_img real_img * (1 - mask) # 训练生成器 fake_img generator(torch.cat([masked_img, mask], dim1)) rec_loss l1_loss(fake_img, real_gt, mask) perc_loss vgg_loss(fake_img, real_gt) adv_loss gan_loss(discriminator(fake_img), real_label) g_loss rec_loss * w_rec perc_loss * w_perc adv_loss * w_adv g_optimizer.zero_grad() g_loss.backward() g_optimizer.step() # 训练判别器 real_pred discriminator(real_gt) fake_pred discriminator(fake_img.detach()) d_loss (gan_loss(real_pred, real_label) gan_loss(fake_pred, fake_label)) * 0.5 d_optimizer.zero_grad() d_loss.backward() d_optimizer.step()这里面有一个很容易忽略的细节计算重建损失时建议把mask也作为权重传入只计算Mask区域内的像素差。如果把全图都算进去背景区域占了绝对主导模型很容易学到“把背景复制一下就行”对被遮挡区域毫无生成能力。4. 实战从零训练与图像修复推理4.1 用自己的数据集跑通训练建议第一次跑用一个小数据集比如从公开数据集里挑500张图片设置训练步数1000步目标不是效果好而是验证整个链路是通的数据加载正确、模型前向正常、损失能下降、权重能保存和加载。这一步通过后再全量训练能省下大量排查时间。使用自己的数据集时先创建一个data_list.txt每一行是图片路径/path/to/train/000001.jpg /path/to/train/000002.jpg ...然后修改config把data_root指向这个txt设置好image_size、batch_size等参数。执行训练python train.py --config configs/train_config.yaml训练启动后观察前几个迭代的日志。正常情况下loss应该在稳步下降PSNR和SSIM缓慢上升。如果loss曲线剧烈震荡我的经验是先把学习率调低一个数量级比如从1e-4调到1e-5再看是否稳定。如果生成器loss和判别器loss走势极不平衡参考上一节提到的调整训练频率。训练时长方面单张RTX 3090跑256×256输入、batch_size810万步大约需要2天。如果想快速验证可以先用--max_iters 2000跑一小段。4.2 推理流程与效果对比推理流程比训练简单但精度问题同样不可忽视。命令类似python infer.py --image samples/test.jpg --mask samples/mask.png --checkpoint output/latest.pth --output result.png推理脚本内部一般做这四步读取图像和Mask统一resize到训练尺寸。图像和Mask拼接经过生成器前向计算。生成结果与原始图像融合保留原始图中未损坏区域。保存输出图像必要时做后处理。这里的关键是融合。不能把生成器的整个输出直接作为结果因为它在非Mask区域也会产生偏移导致原图背景被改。正确的融合方式是result original_image * (1 - mask) generated_image * mask如果发现修复区域边缘有接缝可以对Mask做一次高斯模糊或者把Mask膨胀几个像素让融合过渡更自然。颜色偏色问题时检查推理时是否做了和训练一样的归一化。很多项目训练时把像素映射到[-1,1]推理时忘了减均值除方差出来的图就会明显偏色。效果对比时我习惯把原图、Mask、修复结果拼在一张画布里方便肉眼评估。重点关注三块边缘过渡是否自然、纹理是否重复、颜色是否一致。4.3 评估指标PSNR/SSIM如何看在有真实完整图像的测试集上可以计算PSNR和SSIM。PSNR越高越好通常修复模型能到25~35dBSSIM越接近1越好。另一个更接近主观感受的指标是LPIPS值越低越好。但客观指标和人的感受经常不一致。GAN生成的结果纹理细腻但像素级偏移大PSNR可能反而不如模糊的结果。因此评估时必须同时看主观效果图不能只看分数。我自己的习惯是先用PSNR/SSIM筛掉明显不行的模型再用肉眼对比候选模型的修复图最后根据应用场景做决定。如果业务场景对真实性要求高LPIPS和人工评测优先级要提前。5. 常见问题与排查技巧实录5.1 训练不收敛或损失异常训练时遇到loss变成NaN最先检查三件事输入图像有没有全黑或全白、Mask是否全0、学习率是否过大。图像修复模型输入是全零区域时前向计算容易出现梯度异常。可以先把学习率降到1e-5跑一个iteration看看如果还NaN就要检查数据归一化和模型权重初始化。另一个常见情况是loss快速下降但不代表效果好。生成器可能学会了一种投机取巧的方式非Mask区域直接复制Mask区域输出平均色。这样重建loss很低PSNR也可能不差但视觉上一塌糊涂。解决办法是提高感知损失和对抗损失的权重并定期人工查看生成图。我遇到的训练崩溃还有一个隐藏原因用FP16混合精度训练时loss数值不稳定。如果源码默认开了amp可以关闭再试或者调整grad_scaler的初始化比例。修复模型对精度比较敏感稳定优先于速度。5.2 修复结果模糊或伪影严重模糊是因为模型没有高频细节常见原因包括重建损失权重过高、网络容量不足、输入分辨率过低。我的做法是先降低lambda_rec同时确保lambda_perceptual不是0否则模型会大量丢失纹理。如果还模糊考虑换更大的生成器或增加中间通道数。伪影一般出现在GAN训练不稳定时。表现为修复区域有斑点、彩色条纹或过锐利的边缘。可以考虑降低对抗损失权重。使用LSGAN或Hinge Loss代替标准BCE。判别器增加谱归一化。训练前期冻结判别器只训练生成器若干千步。伪影还可能与推理分辨率有关。如果训练是256×256推理时给512×512的图模型没见过这么大尺寸容易产生结构畸变。此时优先保证推理尺寸和训练一致再考虑是否使用支持任意分辨率的结构。5.3 环境兼容性与显存不足问题环境兼容性问题中最常见的是PyTorch和CUDA版本不匹配。安装PyTorch时如果选了比驱动支持更高的CUDA版本会直接报CUDA driver version is insufficient。解决办法是用nvidia-smi确认支持版本然后重新安装匹配的PyTorch。显存不足的报错形式一般是RuntimeError: CUDA out of memory. Tried to allocate X MiB我推荐的排查顺序降低batch_size到2甚至1。降低图像尺寸比如从512降到256。设置torch.backends.cudnn.benchmarkTrue有时候能省显存。使用梯度累积模拟大batch。开启torch.utils.checkpoint梯度检查点用计算换显存。使用混合精度训练FP16。注意DataLoader的num_workers调大并不会减少显存占用反而可能因为多进程缓存导致整体内存上涨。显存不足时可以先关掉验证集评估因为验证过程也占用显存。5.4 常见问题速查表问题可能原因解决办法loss为NaN学习率过大或输入包含无效值降低学习率检查数据归一化修复区域模糊重建损失权重过高或网络容量不足降低lambda_rec增加网络宽度边界有接缝推理融合未处理Mask对Mask做膨胀或高斯模糊后融合颜色偏色推理归一化与训练不一致统一预处理逻辑CUDA out of memorybatch_size/分辨率过大调小尺寸使用FP16梯度累积预训练权重加载失败模型结构和权重键名不匹配打印keys逐一对比手动修改state_dict大型Mask修复效果差训练时Mask面积分布不匹配增加不规则大Mask的数据增强6. 项目扩展方向与个人心得6.1 如何扩展到自己的业务场景单纯跑通一个开源项目不算本事能把它用起来才是目标。假设你要做老照片修复需要自己准备一批带划痕的图片或者模拟划痕生成训练数据要移除水印就得专门生成文字区域的Mask。千万不要通用模型一把梭效果大概率不如预期。扩展方面可以考虑模型导出。PyTorch模型可以用ONNX导出再用TensorRT做推理加速方便部署到服务端torch.onnx.export( generator, (dummy_img, dummy_mask), inpaint.onnx, input_names[image, mask], output_names[output], dynamic_axes{image: {0: batch}, mask: {0: batch}} )ONNX导出时要注意自定义层比如部分卷积中的Mask更新可能不支持导出需要把部分卷积改写成兼容算子或者用torch.onnx.export时设置custom_opsets。这块我踩过坑建议先导出一个最简模型验证通道。另外把修复系统包装成API接口也很实用。用FastAPI加载模型接收图像和Mask返回修复结果这样就能接入小程序、网页或后端服务。要注意的是并发场景下显存可能不够可以考虑单进程单GPU队列方式避免多个请求同时占用显存导致OOM。6.2 实操后的几点建议最后说几点我反复踩坑之后总结出来的经验。第一个是“先复现再创新”。拿到源码后不要着急改结构先用作者给的预训练权重和测试命令跑通确认整个链路能出成果再开始替换模型、改损失函数否则出现问题根本分不清是原代码的锅还是你自己改的锅。第二个是“每次只改一个变量”。图像修复模型效果受多个因素影响如果一次改了三四个参数效果变差了你完全不知道是哪个参数导致的。我习惯用实验管理工具记录每次训练的配置、log和输出图比如简单地把每个实验放在单独目录命名带上参数摘要。第三个是保存模型时把配置一起保存。只存state_dict有个问题过几个月你自己都忘了当时用了什么image_size、什么损失权重、什么归一化方式。最好是存一个字典包含model_state_dict、config和源码版本方便日后加载和新环境复现。在实际修复效果上别迷信单模型很多时候“检测Mask 修复 后处理”的pipeline比单独调模型更有效。先用分割模型自动生成Mask修复后再用锐化或颜色校正提升观感。这个思路能大幅提高项目的落地价值也是我把这个源码项目吃透之后最想分享的一点。本文还有配套的精品资源点击获取