简介一个基于PyTorch的深度学习图像修复与校准项目聚焦受损或不完整图像的复原适用于数字文化遗产保护、图像增强、影视后期制作等场景面向有一定深度学习基础、希望动手实践图像修复的开发者。资源共18个文件压缩包大小35.81MB包含5个Python脚本数据预处理、模型结构、训练主流程、图像变形辅助、5个PyTorch权重文件、5张效果对比示例图以及LICENSE、README和训练日志等文档目录结构清晰便于查阅。从dataset.py的数据读取到train.py的训练循环再到model.py中的网络定义完整展示了图像修复项目的基本链路借助预训练权重和示例图片可直接对比修复前后效果也能基于已有模型继续微调。此外训练日志记录了损失值等关键信息便于分析模型收敛情况与调优方向。目前已有408人学习适合在PyTorch框架下研究图像修复与校准的开发者参考。1. 图像修复校准先定义“退化”再让网络反推图像修复不是简单的抠图补洞而是基于退化观测反推原始内容的病态求逆问题。不少人把 inpainting 和 restoration 混为一谈前者借助掩码先验填补缺失区域后者需要显式建模模糊、噪声、形变等多类退化过程。这个基于 PyTorch 的项目把两条路收进同一条流水线——deform.py 对原始图施加可控形变生成成对样本model.py 里带跳跃连接的编码器-解码器网络学习重建映射train.py 负责训练并按阶段保存 checkpoint。值得注意的细节是 checkpoint 目录里既有按样本形态区分的 model-emoji01.pt 到 04也有按几何特性划分的 model-shape.pt说明作者的校准思路不是单模型走天下而是按数据形态做模型细分。对想用 PyTorch 入门图像修复又不想只看论文的开发者这个项目的文件组织可以当一份可执行的地图。2. 退化模拟与数据管线从 deform.py 到 dataset.py2.1 退化建模为什么是修复的第一步图像修复任务与普通超分不同它对配对数据的质量非常敏感。如果目标图像是规整的印刷体、字符或界面截图而退化环节只是简单加一层高斯噪声训练出来的模型面对真实场景里的局部扭曲、扫描倾斜、透视压缩时几乎必然失效。常见做法是先固定“退化类别”再为每个类别准备一个可参数化的变形器弹性形变用随机位移场加平滑重采样局部扭曲用仿射或薄板样条TPS模糊退化则用不同核宽的高斯核。deform.py 在这个项目里担任的正是这个角色它把“如何破坏图像”的规则封装成函数dataset.py 在读取每张样本时调用它产出成对的degraded, target训练对。这里有一个容易被忽略的工程点退化算子必须同时满足可复现和足够随机。可复现是为了在训练异常时能回溯到具体某一对样本足够随机是为了防止网络在固定的退化模式上过拟合。项目里如果发生训练损失持续下降但验证集修复效果上不去的现象第一个要检查的就是退化参数的随机范围是否设得太窄导致验证时遇到的破坏模式完全不在训练分布内。2.2 弹性形变的实现与参数解读以下是一版 deform.py 里最常用的弹性形变实现跑通最小可用版本只需要 scipy 的map_coordinates和gaussian_filter# deform.py — 弹性形变核心函数 import numpy as np from scipy.ndimage import gaussian_filter, map_coordinates def elastic_deform(image, alpha25, sigma4, random_state42): 对图像施加平滑的弹性形变返回与输入同尺寸的变形图。 alpha : 最大位移幅度像素值越大破坏越明显 sigma : 位移场高斯平滑尺度控制形变偏全局还是偏局部 random_state : 固定随机源保证同一张图多次处理结果一致。 rng np.random.RandomState(random_state) h, w image.shape[:2] dx rng.uniform(-alpha, alpha, size(h, w)) dy rng.uniform(-alpha, alpha, size(h, w)) dx gaussian_filter(dx, sigma) dy gaussian_filter(dy, sigma) rows, cols np.meshgrid(np.arange(h), np.arange(w), indexingij) coords (rows dy, cols dx) if image.ndim 2: return map_coordinates(image, coords, order1, modereflect) deformed np.stack( [map_coordinates(image[..., c], coords, order1, modereflect) for c in range(image.shape[-1])], axis-1 ) return deformed这段代码的逻辑是用随机噪声生成二维位移场再用高斯滤波把位移场变平滑最后按新坐标对原图重采样。alpha 决定退化的强度上限建议前几个 epoch 用较小的 alpha 让网络先学低频结构再逐步加大到目标值sigma 保持在 3 到 6 之间比较合适过大会让位移场过于全局局部畸变退化成整体平移过小则会出现撕裂感。order1 表示双线性重采样比 order0 的最近邻平滑也不像 order3 那样引入振铃伪影。modereflect 保证边界像素不会把空值带进训练张量这一步在 PyTorch 环境里经常被忽略一旦边界填充缺失训练 loss 会在每轮开头出现一个不合理的尖峰。2.3 dataset.py 的构造细节与常见误用Dataset 类的骨架如下需要特别关注__getitem__里退化处理必须发生在 tensor 化之前否则 numpy 中转会断开 PyTorch 的自动求导链路# dataset.py — 修复任务的数据集封装 import os import numpy as np from PIL import Image from torch.utils.data import Dataset from torchvision.transforms import ToTensor class RestorationDataset(Dataset): def __init__(self, root_dir, deform_fnNone, size(256, 256)): self.paths [os.path.join(root_dir, f) for f in os.listdir(root_dir) if f.lower().endswith((.png, .jpg, .jpeg))] self.deform_fn deform_fn self.size size self.to_tensor ToTensor() def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) img img.resize(self.size) if self.deform_fn is not None: # 退化处理必须在归一化前完成保持 0~255 数值域统一 degraded self.deform_fn(np.array(img)) degraded self.to_tensor(degraded.astype(np.float32) / 255.) target self.to_tensor(np.array(img).astype(np.float32) / 255.) return degraded, target # 无退化函数时返回同一张图的双份拷贝便于自监督对比实验 arr np.array(img).astype(np.float32) / 255. return self.to_tensor(arr), self.to_tensor(arr)退化函数的输入输出要保持在同一个数值范围。若变形前是 0 到 255变形后也应该是 0 到 255不要在中间混用 0 到 1 和 0 到 255否则训练初期 loss 会出现剧烈震荡。DataLoader 的 num_workers 建议设到 2 以上数据增强的随机种子挂在 worker_init_fn 上这样即使在多进程读取下每一张图对应的退化结果仍然可以按种子回溯排错时能精确锁定是数据问题还是模型问题。实际项目中我曾见过__getitem__里反复调用 PIL resize 导致训练吞吐骤降的情况正确做法是在__init__阶段统一做一次 resize__getitem__只做增强、变形和转 tensor。3. 网络架构model.py 的编码器-解码器与上下文融合3.1 修复任务为什么需要编码器-解码器结构图像修复要求输出像素级的预测且输入输出空间尺寸完全一致。普通分类网络通过全局池化把空间信息压扁最终输出类别概率无法用于像素级重建。编码器-解码器结构的价值在于编码器逐级降低分辨率每下采样一次感受野扩大一倍让网络看到更大的上下文。修复一个被划伤的文字区域看不到周围几行文字的结构不可能补得合理。解码器逐级上采样恢复细节同时通过跳跃连接把编码器同层的结构信息直传过来避免空间细节在多次降采样中彻底丢失。这个项目里 checkpoint 命名出现 emoji 和 shape 两类标签合理的推测是作者针对不同的训练数据分布分别训练了多套权重。如果训练集里既有整页文档图像也有小块图标单模型要同时处理大尺度形变和小范围纹理修补网络容量会明显紧张。更稳妥的做法是为每一类数据单独训练一个模型推理时通过一个分类前置模块做路由这也是 model-emoji.pt 和 model-shape.pt 可以并存的原因。3.2 跳跃连接与可变形卷积的取舍项目文件里 deforme.py 与 model.py 同时存在后者定义主干网络前者负责形变相关的辅助模块。可变形卷积与普通卷积的核心区别在于卷积核的采样位置可以随输入内容偏移不固定在规则网格上。在图像修复里这条性质很有用当退化包含弹性形变或投影畸变时普通卷积只能靠堆叠层数间接建模几何变化可变形卷积第一层就能把采样网格偏移到退化发生的位置等效于提前做了一个可学习的空间变换。代价是梯度计算比普通卷积重约 30%小分辨率验证集上收益明显大图训练时需要留意显存峰值。跳跃连接的数量和位置则直接影响修复质量。只拼接最深层特征语义信息足够但高频细节缺失逐层拼接参数和激活值占用都会明显上升。我一般保留两层跳跃连接分别在最低分辨率和中间层位置最后一层 head 前加 1x1 卷积做通道压缩减小输出特征图上的冗余通道。3.3 前向传播的代码骨架以下是 model.py 中可复用的一版小型 U-Net 风格实现# model.py — 编码器-解码器修复网络骨架 import torch import torch.nn as nn class DownBlock(nn.Module): 编码器模块stride2 卷积完成下采样同时翻倍通道数。 def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size3, stride2, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class UpBlock(nn.Module): 解码器模块转置卷积上采样后拼接跳跃连接。 def __init__(self, in_ch, out_ch): super().__init__() self.up nn.ConvTranspose2d(in_ch, out_ch, kernel_size2, stride2) self.conv nn.Sequential( nn.Conv2d(out_ch * 2, out_ch, kernel_size3, padding1), nn.ReLU(inplaceTrue), ) def forward(self, x, skip): x self.up(x) return self.conv(torch.cat([x, skip], dim1)) class RestorationNet(nn.Module): def __init__(self, in_ch3, base_ch32): super().__init__() self.enc1 DownBlock(in_ch, base_ch) self.enc2 DownBlock(base_ch, base_ch * 2) self.enc3 DownBlock(base_ch * 2, base_ch * 4) self.dec2 UpBlock(base_ch * 4, base_ch * 2) self.dec1 UpBlock(base_ch * 2, base_ch) self.head nn.Conv2d(base_ch, in_ch, kernel_size3, padding1) def forward(self, x): e1 self.enc1(x) # H/2, base_ch e2 self.enc2(e1) # H/4, base_ch*2 e3 self.enc3(e2) # H/8, base_ch*4 d2 self.dec2(e3, e2) # 与中间层特征拼接 d1 self.dec1(d2, e1) # 与原始尺度细节拼接 return self.head(d1)forward 里两个拼接点是网络的关键设计位置。e3 向下传递语义信息e2 和 e1 分别保留不同粒度的边缘与纹理。head 层没有接激活函数这便于在训练中使用带有负区间的损失组合。base_ch 从 16 到 64 都可以但这是一个 200 个 epoch 就能收敛的规模区间再往上加通道训练显著变慢而修复精度提升有限。若输入尺寸固定为 256x256编码器下采样两次后特征图为 64x64显存占用比较可控若输入尺寸不固定BatchNorm 的均值和方差会在 batch 内波动建议改用 GroupNorm 或固定所有训练样本的尺寸。4. train.py 训练循环与 checkpoint 恢复机制4.1 训练循环的完整框架训练脚本在 PyTorch 项目里套路最固定也最容易出现“看着正常、跑起来异常”的问题。一个完整的训练循环需要覆盖模型初始化、优化器与调度器设置、断点恢复、前向传播、梯度裁剪、混合精度更新这几个环节# train.py — 训练主循环节选 import os import torch import torch.nn as nn from torch.utils.data import DataLoader device torch.device(cuda if torch.cuda.is_available() else cpu) model RestorationNet(in_ch3).to(device) criterion nn.L1Loss() # 图像修复中 L1 比 MSE 收敛更稳 optimizer torch.optim.AdamW(model.parameters(), lr2e-4, weight_decay1e-5) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max200) scaler torch.cuda.amp.GradScaler() # 混合精度训练 start_epoch 0 ckpt_path checkpoint/model-emoji03.pt if os.path.exists(ckpt_path): ckpt torch.load(ckpt_path, map_locationdevice) model.load_state_dict(ckpt[model_state_dict]) optimizer.load_state_dict(ckpt[optimizer_state_dict]) start_epoch ckpt[epoch] 1 for epoch in range(start_epoch, 200): model.train() for degraded, target in train_loader: degraded, target degraded.to(device), target.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): pred model(degraded) loss criterion(pred, target) scaler.scale(loss).backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5) scaler.step(optimizer) scaler.update() scheduler.step()这段代码的逻辑是每个 batch 内清零梯度、混合精度前向计算损失、反向后裁剪梯度再更新参数。AdamW 的 weight_decay 与 Adam 不同它是解耦的权重衰减图像修复类的像素回归任务里比 Adam 更稳。clip_grad_norm 的 0.5 是保守值能有效防止由于退化样本剧烈变化引发的梯度爆炸。混合精度训练时损失计算必须在autocast()上下文里完成GradScaler 负责把缩放后的梯度恢复到正常数值范围后再更新权重如果训练过程中 loss 出现 NaN优先检查这一步是否被遗漏。4.2 损失函数与梯度策略项目日志文件 logs.txt 里如果出现 loss 断崖式下降又快速回升往往不是学习率问题而是损失函数对异常像素的敏感度设置不对。L1 Loss 对离群点不敏感适合大多数修复任务Charbonnier Loss 是 L1 平滑近似收敛更稳如果加入感知损失则需要一个预训练分类网络提取中间层特征计算量会增加一截。对于 checkpoint 命名里出现多个类似语义的模型我的做法是让主损失用 L1附加 0.05 权重的感知损失并在验证集上按 PSNR 和 SSIM 两个指标决定是否覆盖保存最佳权重。全变分正则TV Loss只在输出出现明显棋盘效应时加入正则权重建议控制在 1e-4 以下。4.3 checkpoint 保存策略与断点续训checkpoint 目录里的 model-emoji.pt 到 model-emoji04.pt 说明作者在按训练阶段滚动保存这种习惯在长实验里非常值得保留。不同存储方式的选择需要根据场景决定存储格式包含内容体积适用场景state_dict 单权重仅模型参数最小推理部署不关心训练状态完整 checkpointmodel optimizer scheduler epoch适中断点续训需要恢复完整状态全量序列化模型结构 权重 Python 对象最大快速恢复实验环境但跨 PyTorch 版本兼容性差断点保存代码里需要把 epoch、模型权重、优化器状态、最佳验证指标一并写入否则加载后只能恢复参数无法恢复学习率调度位置。PyTorch 2.x 之后torch.load默认weights_onlyTrue如果 checkpoint 里保存了自定义 Python 对象加载时会直接报错这时需要显式传weights_onlyFalse或者只保存纯 tensor 和 dict 结构保持向前兼容。多 GPU 训练保存的权重在加载时会多出module.前缀解决办法是保存时去掉前缀或加载时做一次字符串替换这个问题在推理阶段经常碰到。5. 从 checkpoint 到推理结果加载、评估与常见坑5.1 加载模型并跑通第一次推理拿到 checkpoint 目录里某个 .pt 文件后先确认它保存的是 state_dict 还是完整 checkpoint然后用下面的代码加载推理# inference.py — 从 checkpoint 加载模型并推理 import torch from torchvision.transforms import ToTensor, ToPILImage device torch.device(cuda if torch.cuda.is_available() else cpu) model RestorationNet(in_ch3).to(device) ckpt torch.load(checkpoint/model-shape.pt, map_locationdevice, weights_onlyFalse) model.load_state_dict(ckpt.get(model_state_dict, ckpt)) model.eval() # 关键关闭 dropout 和 batchnorm 的 batch 统计 degraded ToTensor()(Image.open(images/deformed_img.png).convert(RGB)) degraded degraded.unsqueeze(0).to(device) with torch.no_grad(): pred model(degraded) result ToPILImage()(pred.squeeze(0).cpu().clamp(0, 1)) result.save(images/prediction_results.png)加载路径上的 map_location 负责把 GPU 权重映射到当前可用设备。get(model_state_dict, ckpt)用于兼容两种存储方式完整 checkpoint 取 model_state_dict 键纯 state_dict 直接作为权重字典。推理时若已保存的模型由多卡训练得到权重键名带module.前缀在 load_state_dict 前需要加上collections.OrderedDict((k.replace(module., ), v) ...)的替换逻辑。最后结果 clamp 到 0 到 1 再保存避免因网络输出超出数值范围导致图像整体偏灰或出现反白。5.2 训练推理不一致与模型加载的典型故障最常见的推理异常是训练时效果正常、推理时输出偏模糊。这通常不是模型问题而是model.eval()没有调用。训练模式下 BatchNorm 使用当前 batch 的均值和方差推理模式下应该使用累计的全局统计量不切换的话batch 为 1 时输出会明显失真。另一个高频问题是 torch.load 在 PyTorch 2.6 之后对新保存的权重自动启用 weights_only 限制如果 checkpoint 里嵌入了自定义类加载会直接抛出反序列化异常处理方式是保存时拆成纯 dict 结构或加载时显式关闭限制。最后修复输出的尺寸必须和训练时保持一致训练用 256x256推理时换成 512x512padding 行为会改变修复区域边缘会出现一圈纹理突变出现这类情况时先检查输入尺寸再检查归一化均值方差是否和训练管线一致。本文还有配套的精品资源点击获取