简介本资源是一套基于SwinIR模型的自定义图像恢复训练与测试代码实现面向计算机视觉方向的研究者与深度学习初学者聚焦图像超分辨率重建与去噪等典型低级视觉任务。代码逻辑完整、注释清晰支持开箱即用图像去噪任务仅需调整数据路径即可运行超分任务需在数据加载类中取消patchsize相关操作便于理解底层数据处理机制。压缩包共17个文件含5个核心Python源码如main.py、net.py、data.py、4个编译缓存pyc、4个IDE配置XML文件及MATLAB评估脚本psnr.m、compute_psnr.m等整体仅28KB轻量易读。已有5370人学习下载代码结构层次分明涵盖数据集组织inputTrain/targetTrain等目录、网络构建、损失定义、训练循环与结果评估全流程特别适合复现SwinTransformer在图像恢复中的应用并开展二次开发。1. 项目概述从“拿来就用”到“深度定制”的SwinIR训练之旅如果你在图像超分辨率、去噪或者JPEG压缩伪影去除这些任务上折腾过大概率听说过SwinIR这个名字。它基于Swin Transformer架构在多个图像复原的公开基准测试上都刷出了非常亮眼的成绩。很多朋友拿到官方代码仓库跑通预训练模型在标准数据集上的测试感觉效果不错但一到想用自己的数据、调整网络结构或者尝试新的损失函数时就卡住了。官方代码库为了通用性和可复现性往往结构复杂各种配置文件和抽象层叠在一起对于想快速验证一个想法或者进行教学演示的人来说学习成本太高。这就是我动手整理这份“逻辑完整易懂”的自定义训练测试代码的初衷——剥离所有不必要的抽象把数据流、模型前向传播、损失计算、优化器更新、验证测试这几个核心环节用最直白、模块化的Python代码呈现出来。让你不仅能跑起来更能看清楚每一行代码在做什么知道在哪里插入你自己的修改。这份代码的核心价值在于“透明”和“可插拔”。它不追求大而全的工程化而是聚焦于SwinIR模型训练与测试最本质的流程。你将看到数据是如何被加载、预处理并送入模型的损失函数是如何计算并反向传播的模型权重是如何保存和加载的。更重要的是代码结构清晰每个功能块数据加载、模型定义、训练循环、测试评估都是独立的你可以像搭积木一样替换其中的任何部分。例如想把L1损失换成感知损失Perceptual Loss只需要修改训练循环中的几行代码。想尝试不同的学习率调度策略替换对应的模块即可。这对于研究者快速进行算法迭代或者初学者深入理解深度学习训练流程都是非常有帮助的。2. 核心模块拆解你的数据、模型与训练循环一份完整的训练代码可以抽象为三个核心部分数据管道、模型定义和训练引擎。下面我们逐一拆解并解释在SwinIR这个具体场景下每一部分是如何实现的。2.1 数据加载器打造专属的图像对流水线图像复原任务如超分的训练数据通常是“低质量-高质量”图像对。我们的数据加载器Dataset和DataLoader需要高效地读取、配对并预处理这些图像。2.1.1 自定义Dataset类的设计逻辑我们首先定义一个SRDataset类。它的核心工作是给定一个包含低分辨率LR图像和高分辨率HR图像的文件夹路径在初始化时建立图像对的路径列表。这里的关键设计点是如何组织你的数据。一个常见的简单结构是your_dataset/ ├── LR/ # 存放所有低分辨率图像 │ ├── 001.png │ ├── 002.png │ └── ... └── HR/ # 存放所有对应的高分辨率图像 ├── 001.png ├── 002.png └── ...要求LR和HR文件夹中的文件名严格一一对应。SRDataset的__getitem__方法会接收一个索引然后同时读取对应索引的LR和HR图像。2.1.2 图像预处理与增强的关键步骤读取图像通常使用PIL或OpenCV后不能直接丢给模型。需要经过一系列预处理变换这些变换通过PyTorch的torchvision.transforms组合成一个transform管道。对于训练阶段典型的流程包括随机裁剪为了增加数据的多样性和实现批训练我们从HR图像中随机裁剪一个固定大小的块如patch_size128然后在LR图像对应的位置考虑缩放因子例如超分4倍则LR中裁剪128/432的块进行同步裁剪。这确保了图像对的区域一致性。随机旋转与翻转这是最常用且几乎无成本的几何数据增强。以一定概率对图像对进行90°、180°、270°的随机旋转以及水平/垂直翻转可以极大增加数据多样性防止模型过拟合。转换为张量将PIL图像或NumPy数组转换为PyTorch张量并将像素值范围从[0, 255]归一化到[0.0, 1.0]。可选通道归一化有时会进一步根据数据集的均值和标准差进行归一化但对于图像复原使用[0,1]范围通常已足够。对于测试/验证阶段变换则简单得多通常只需要将图像缩放到统一大小或保持原尺寸然后转换为张量和归一化。一个重要的细节测试时如果图像尺寸不能被模型的缩放因子整除可能会在边界产生伪影。因此有时需要对测试图像进行填充padding以满足整除条件在输出后再将填充部分裁剪掉。2.1.3 DataLoader的参数调优经验将定义好的Dataset交给DataLoader。有几个参数直接影响训练效率batch_size根据你的GPU显存决定。SwinIR模型不算小在消费级显卡如RTX 4090上超分任务可能batch_size只能设为8或16。可以从一个小值开始尝试逐步增加直到显存占满。num_workers数据加载的并行进程数。建议设置为CPU核心数或略少可以显著加速数据加载避免训练循环等待数据。在Linux系统上效果尤为明显。pin_memoryTrue当使用GPU时将此参数设为True可以将数据更快地从主机内存传输到GPU显存提升训练速度。shuffleTrue仅在训练时使用打乱数据顺序使每个epoch看到的数据分布都不同。注意如果你的数据集非常大例如数万对图像且存储于机械硬盘过高的num_workers可能会导致磁盘I/O成为瓶颈。此时需要平衡。使用SSD硬盘可以极大缓解此问题。2.2 模型定义理解SwinIR的骨干与头SwinIR的官方实现通常包含多个部分浅层特征提取、深层特征提取Swin Transformer块和高质量图像重建。在我们的简化版代码中我们会聚焦于核心部分。2.2.1 浅层特征提取从像素到特征这部分通常是一个简单的卷积层nn.Conv2d。它的作用是将输入的3通道RGB低质量图像映射到一个更高维的特征空间例如64或128个通道。你可以把它理解为一个“编码器”的入口它将原始的像素信息转换为一系列更抽象、更适合后续深度网络处理的特征图。这个卷积层的核大小通常是3x3步长为1填充为1以保持空间分辨率不变。2.2.2 核心骨干Swin Transformer块堆叠这是SwinIR的灵魂。浅层特征会被送入一系列Swin Transformer块Swin Transformer Blocks中。每个块内部主要包含窗口多头自注意力W-MSA将特征图划分为不重叠的局部窗口如8x8在每个窗口内计算自注意力。这大大降低了传统全局自注意力的计算复杂度从图像尺寸的平方级降低到窗口尺寸的平方级。移位窗口多头自注意力SW-MSA为了引入窗口间的信息交互在连续的Swin Transformer块中会交替使用W-MSA和SW-MSA。SW-MSA会将窗口进行一定偏移例如向右下角偏移半个窗口使得当前窗口能与上一层的不同窗口区域进行交互。多层感知机MLP在每个注意力层之后会接一个两层的全连接网络中间带有GELU激活函数用于进一步的特征变换。层归一化LayerNorm和残差连接每个子层注意力、MLP前都应用层归一化并广泛使用残差连接来缓解深层网络的梯度消失问题。在代码中这部分通常被组织为一个nn.ModuleList里面包含多个RSTBResidual Swin Transformer Block或类似命名的模块。一个关键的超参数是“深度”即堆叠多少个这样的块。更多的块意味着更大的模型容量和更长的训练时间但也可能带来性能提升需要根据任务和数据集权衡。2.2.3 重建模块从特征回到图像经过深层特征提取后我们得到了富含上下文信息的特征图。重建模块的任务是将这些特征图“翻译”回高质量的RGB图像。对于超分任务这通常是一个亚像素卷积层nn.PixelShuffle的上采样方式或一个转置卷积层nn.ConvTranspose2d后接一个最终的卷积层将通道数调整回3。对于去噪或去块效应任务可能只需要一个卷积层来直接预测残差即“干净图像噪声图像预测的噪声”或重建图像。2.2.4 预训练权重的加载与部分微调SwinIR官方提供了在多个大型数据集如DIV2K上预训练的模型。在我们的代码中加载预训练权重是一个非常重要的步骤可以让你在小数据集上快速获得好效果。使用model.load_state_dict(torch.load(pretrain_path), strictFalse)来加载。strictFalse参数很关键因为它允许只加载模型和预训练权重中名称匹配的部分对于你修改过的层例如改变了输出通道数可以忽略从而支持部分微调。如果你想冻结骨干网络只训练重建头可以这样做for name, param in model.named_parameters(): if ‘head’ not in name: # 假设重建层的参数名包含‘head’ param.requires_grad False这样在反向传播时只有requires_gradTrue的参数才会更新可以节省显存并加速训练初期收敛。2.3 训练循环损失、优化与评估的舞蹈这是整个代码的“发动机”一个标准的PyTorch训练循环但其中每一步都有讲究。2.3.1 损失函数的选择L1还是L2图像复原任务最常用的损失函数是L1损失MAE和L2损失MSE。L2损失MSE惩罚大的误差更重倾向于产生更平滑的输出但有时会导致图像过于模糊丢失高频纹理细节。L1损失MAE对误差的惩罚是线性的通常能产生更清晰的边缘和纹理是目前超分等任务的主流选择。在我们的代码中你可以简单地使用nn.L1Loss()。但更优的做法是结合多种损失。例如感知损失Perceptual Loss利用预训练的分类网络如VGG提取特征在特征空间计算差异使复原图像在“感知”上更接近真实图像提升视觉质量。对抗损失GAN Loss引入一个判别器网络让生成器SwinIR产生尽可能“以假乱真”的图像可以极大地提升纹理的真实感。但这会使得训练变得不稳定需要仔细调整。在简化版代码中我们可能只使用L1损失作为起点。你可以通过以下方式轻松扩展criterion_pixel nn.L1Loss() criterion_perceptual PerceptualLoss() # 需要自定义或引用第三方库 ... loss criterion_pixel(output, target) 0.01 * criterion_perceptual(output, target)2.3.2 优化器与学习率调度寻找下降的路径对于SwinIR这类Transformer-based模型Adam或AdamW优化器是标准选择。AdamW相比Adam加入了权重衰减的正则化通常泛化性能更好。optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4)初始学习率lr是一个关键超参数。对于微调预训练模型通常使用较小的学习率如1e-5到1e-4对于从头训练可以稍大如1e-4。学习率调度同样重要。最常见的是余弦退火调度CosineAnnealingLR它让学习率随着训练过程从初始值平滑地衰减到0模拟了“精细调参”的过程。另一种实用的策略是带热重启的余弦退火CosineAnnealingWarmRestarts它周期性地重启学习率有助于模型跳出局部最优。在我们的代码框架中你可以很方便地集成这些调度器。2.3.3 训练迭代中的关键操作一个训练迭代iteration包含以下步骤清零梯度optimizer.zero_grad()。这是必须的因为PyTorch会累积梯度。前向传播outputs model(lr_images)。计算损失loss criterion(outputs, hr_images)。反向传播loss.backward()。计算所有可训练参数关于损失的梯度。梯度裁剪可选但推荐torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。这对于稳定Transformer模型的训练尤其重要可以防止梯度爆炸。参数更新optimizer.step()。根据梯度和优化算法更新模型权重。学习率更新scheduler.step()在每个epoch或iteration后取决于调度器类型。2.3.4 验证与模型保存避免过拟合的眼睛我们通常在每个epoch结束后在独立的验证集上评估模型性能。验证时必须将模型设置为评估模式model.eval()这会关闭Dropout、BatchNorm等层在训练和评估时的不同行为。同时要使用torch.no_grad()上下文管理器来禁用梯度计算节省内存和计算资源。评估指标除了损失值还应包括感知指标例如PSNR峰值信噪比最常用的客观指标值越高越好但与主观视觉质量有时不完全一致。SSIM结构相似性比PSNR更能反映人眼对结构信息的感知。模型保存策略通常有两种定期保存每N个epoch保存一次用于检查点恢复。保存最佳模型根据验证集上的PSNR或SSIM只保存指标最好的那个模型权重。这是最常用的策略。if current_psnr best_psnr: best_psnr current_psnr torch.save(model.state_dict(), ‘best_model.pth’)3. 测试与推理让模型真正工作起来训练好的模型最终要用于处理新的图像。测试脚本需要独立、高效且易于使用。3.1 单图与批处理推理脚本推理脚本的核心流程是加载模型权重 - 预处理输入图像 - 模型前向传播 - 后处理输出图像 - 保存。3.1.1 图像预处理对齐训练时配置这一点至关重要测试时对输入图像的预处理裁剪、归一化等必须与训练时完全一致。如果训练时输入是[0,1]范围测试时读入的[0,255]图像就必须除以255。如果训练时对数据进行了减均值除标准差的操作测试时也必须用相同的均值和标准差。不一致的预处理会导致模型性能严重下降甚至产生毫无意义的结果。3.1.2 处理任意尺寸的输入训练时我们使用固定大小的裁剪块但测试时图像尺寸千变万化。SwinIR模型由于其中的Patch Merging/Embedding等操作可能对输入尺寸有整除要求例如需要是窗口大小的整数倍。一个常见的技巧是反射填充import torch.nn.functional as F mod_pad_h (window_size - h % window_size) % window_size mod_pad_w (window_size - w % window_size) % window_size img F.pad(img, (0, mod_pad_w, 0, mod_pad_h), ‘reflect’)在模型输出后再将填充的部分裁剪掉得到最终结果。3.1.3 批处理与GPU内存管理当需要处理大量图像或视频序列时批处理可以提升效率。但要注意大尺寸图像进行批处理会消耗巨量显存。一个稳健的策略是动态调整批处理大小。可以先尝试一个较小的batch_size如1或2如果GPU内存允许再逐步增加。也可以编写一个循环每次处理一批并适时清空缓存 (torch.cuda.empty_cache())。3.2 性能评估与可视化对比单纯的PSNR/SSIM数字有时不够直观。一个好的测试脚本应该能生成可视化的对比图。3.2.1 生成对比网格图将输入的低质量图像、模型输出的复原图像、真实的高质量图像如果有的话并排放在一起保存为一张图片。这能最直观地展示模型的效果。可以使用Matplotlib或OpenCV来轻松实现。更进一步可以计算并标注出每张结果图与GT之间的PSNR/SSIM值。3.2.2 局部区域放大对比人眼对纹理和边缘细节非常敏感。全局图看起来可能都不错但放大到某个特定区域如文字、毛发、建筑纹理差异就显现了。在生成对比图时可以额外插入一个放大特定区域的子图让优劣一目了然。3.2.3 指标批量计算与报告对于整个测试集脚本应该能自动遍历所有图像计算平均PSNR、平均SSIM并生成一个简洁的文本报告或CSV文件。这对于量化比较不同模型或不同训练轮次的效果至关重要。4. 代码实战从零搭建可运行的训练流程让我们抛开复杂的抽象用最直接的代码勾勒出骨架。这里会省略一些非常具体的Swin Transformer块实现细节你可以从官方仓库引用但会完整展示所有接口和流程。4.1 数据准备与Dataset类实现首先确保你的数据按前述结构整理好。然后实现SRDatasetimport os from PIL import Image import torch from torch.utils.data import Dataset import torchvision.transforms as transforms class SRDataset(Dataset): def __init__(self, lr_dir, hr_dir, patch_size128, scale4, is_trainTrue): self.lr_dir lr_dir self.hr_dir hr_dir self.patch_size patch_size self.scale scale self.is_train is_train # 获取所有配对的文件名假设文件名相同 self.lr_filenames sorted([f for f in os.listdir(lr_dir) if f.endswith(‘.png’)]) self.hr_filenames sorted([f for f in os.listdir(hr_dir) if f.endswith(‘.png’)]) # 简单检查 assert len(self.lr_filenames) len(self.hr_filenames) for lr, hr in zip(self.lr_filenames, self.hr_filenames): assert lr hr, f“文件名不匹配: {lr} vs {hr}” # 定义变换 if self.is_train: self.transform transforms.Compose([ transforms.RandomCrop(patch_size), # 对HR进行随机裁剪 transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(90), transforms.ToTensor(), # 自动转换到[0,1] ]) else: self.transform transforms.Compose([ transforms.ToTensor(), ]) def __len__(self): return len(self.lr_filenames) def __getitem__(self, idx): lr_path os.path.join(self.lr_dir, self.lr_filenames[idx]) hr_path os.path.join(self.hr_dir, self.hr_filenames[idx]) lr_img Image.open(lr_path).convert(‘RGB’) hr_img Image.open(hr_path).convert(‘RGB’) # 注意对于训练我们需要同步裁剪LR和HR if self.is_train: # 对HR进行随机裁剪 i, j, h, w transforms.RandomCrop.get_params(hr_img, output_size(self.patch_size, self.patch_size)) hr_img_crop transforms.functional.crop(hr_img, i, j, h, w) # 对LR在对应位置进行裁剪考虑缩放因子 lr_img_crop transforms.functional.crop(lr_img, i//self.scale, j//self.scale, h//self.scale, w//self.scale) # 应用其他增强翻转、旋转 seed torch.randint(0, 2**32, (1,)).item() torch.manual_seed(seed) lr_img_crop self.transform(lr_img_crop) torch.manual_seed(seed) # 确保LR和HR应用完全相同的随机变换 hr_img_crop self.transform(hr_img_crop) return lr_img_crop, hr_img_crop else: lr_img self.transform(lr_img) hr_img self.transform(hr_img) return lr_img, hr_img, self.lr_filenames[idx] # 测试时返回文件名用于保存结果4.2 简化版SwinIR模型定义与训练循环这里我们定义一个极简的模型外壳重点展示训练循环的逻辑。假设我们已经有一个叫SwinIR的模型类可以从官方实现中导入或简化实现。import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm # 用于显示进度条 import time # 假设我们已经定义或导入了SwinIR模型 # from models.network_swinir import SwinIR class SimpleSwinIR(nn.Module): # 此处应为具体的模型定义这里用占位符 def __init__(self, upscale4): super().__init__() # 这里应包含浅层特征提取、Swin Transformer块、重建头等 self.conv_first nn.Conv2d(3, 64, 3, 1, 1) # ... 更多层定义 self.recon_head nn.Conv2d(64, 3, 3, 1, 1) def forward(self, x): # ... 前向传播逻辑 return self.recon_head(x) def train_one_epoch(model, dataloader, criterion, optimizer, device, epoch): model.train() running_loss 0.0 progress_bar tqdm(dataloader, descf‘Epoch [{epoch}]‘) for batch_idx, (lr_imgs, hr_imgs) in enumerate(progress_bar): lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) # 清零梯度 optimizer.zero_grad() # 前向传播 outputs model(lr_imgs) # 计算损失 loss criterion(outputs, hr_imgs) # 反向传播 loss.backward() # 梯度裁剪可选但推荐 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 更新参数 optimizer.step() running_loss loss.item() progress_bar.set_postfix({‘loss’: loss.item()}) avg_loss running_loss / len(dataloader) return avg_loss def validate(model, dataloader, criterion, device): model.eval() running_loss 0.0 total_psnr 0.0 with torch.no_grad(): for lr_imgs, hr_imgs in dataloader: lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) outputs model(lr_imgs) loss criterion(outputs, hr_imgs) running_loss loss.item() # 计算PSNR (假设图像范围在[0,1]) mse torch.mean((outputs - hr_imgs) ** 2) psnr 20 * torch.log10(1.0 / torch.sqrt(mse)) total_psnr psnr.item() avg_loss running_loss / len(dataloader) avg_psnr total_psnr / len(dataloader) return avg_loss, avg_psnr def main(): # 配置参数 device torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’) lr_dir ‘path/to/your/LR/train’ hr_dir ‘path/to/your/HR/train’ val_lr_dir ‘path/to/your/LR/val’ val_hr_dir ‘path/to/your/HR/val’ batch_size 16 num_epochs 100 learning_rate 1e-4 # 数据加载 train_dataset SRDataset(lr_dir, hr_dir, patch_size128, scale4, is_trainTrue) val_dataset SRDataset(val_lr_dir, val_hr_dir, patch_size128, scale4, is_trainFalse) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size1, shuffleFalse, num_workers2, pin_memoryTrue) # 模型、损失、优化器 model SimpleSwinIR(upscale4).to(device) criterion nn.L1Loss() optimizer optim.AdamW(model.parameters(), lrlearning_rate, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxnum_epochs) best_psnr 0.0 for epoch in range(1, num_epochs1): # 训练一个epoch train_loss train_one_epoch(model, train_loader, criterion, optimizer, device, epoch) # 验证 val_loss, val_psnr validate(model, val_loader, criterion, device) print(f‘Epoch {epoch}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val PSNR: {val_psnr:.2f}dB’) # 学习率调度 scheduler.step() # 保存最佳模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(model.state_dict(), f‘best_model_epoch{epoch}_psnr{val_psnr:.2f}.pth’) print(f‘模型已保存 (PSNR: {val_psnr:.2f}dB)’) # 定期保存检查点 if epoch % 10 0: torch.save({ ‘epoch’: epoch, ‘model_state_dict’: model.state_dict(), ‘optimizer_state_dict’: optimizer.state_dict(), ‘scheduler_state_dict’: scheduler.state_dict(), ‘best_psnr’: best_psnr, }, f‘checkpoint_epoch{epoch}.pth’) if __name__ ‘__main__’: main()4.3 测试脚本示例训练完成后使用以下脚本进行单张图像测试import torch from PIL import Image import torchvision.transforms as transforms import numpy as np import cv2 def load_model(model_path, device): model SimpleSwinIR(upscale4).to(device) state_dict torch.load(model_path, map_locationdevice) model.load_state_dict(state_dict) model.eval() return model def preprocess_image(image_path, scale4, window_size8): “”“预处理单张图像包括填充以满足窗口大小整除要求”“” img Image.open(image_path).convert(‘RGB’) img_tensor transforms.ToTensor()(img).unsqueeze(0) # [1, C, H, W] # 填充 _, _, h, w img_tensor.shape mod_pad_h (window_size - h % window_size) % window_size mod_pad_w (window_size - w % window_size) % window_size img_tensor torch.nn.functional.pad(img_tensor, (0, mod_pad_w, 0, mod_pad_h), ‘reflect’) return img_tensor, h, w def test_single_image(model, lr_image_path, output_path, device, scale4): model.eval() with torch.no_grad(): input_tensor, orig_h, orig_w preprocess_image(lr_image_path, scale) input_tensor input_tensor.to(device) output_tensor model(input_tensor) # 裁剪掉填充的部分 output_tensor output_tensor[:, :, :orig_h*scale, :orig_w*scale] # 后处理张量转回图像 output_np output_tensor.squeeze(0).cpu().numpy() # [C, H, W] output_np np.transpose(output_np, (1, 2, 0)) # [H, W, C] output_np np.clip(output_np * 255.0, 0, 255).astype(np.uint8) # 保存 cv2.imwrite(output_path, cv2.cvtColor(output_np, cv2.COLOR_RGB2BGR)) print(f‘结果已保存至: {output_path}’) # 使用示例 device torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’) model load_model(‘best_model.pth’, device) test_single_image(model, ‘test_input.png’, ‘test_output.png’, device)5. 避坑指南与性能调优实战心得在实际操作中你会遇到各种各样的问题。下面是我在多次训练SwinIR和类似模型时总结的一些关键经验和常见陷阱。5.1 训练不收敛或效果差的排查清单如果你的模型训练后效果很差甚至不收敛请按以下顺序检查数据与预处理这是最常见的问题源。首先确保你的LR-HR图像对是严格对齐的。用眼睛看几对放大检查边缘是否对齐。其次检查预处理代码确保训练和测试时的归一化方式是[0,1]还是[0,255]是否减了均值完全一致。一个快速验证方法是取一张训练图像经过Dataset的__getitem__处理后再保存出来看看是否还是正常的图像。损失函数值观察训练初期的损失值。如果损失从一开始就是NaN或无限大很可能是梯度爆炸。尝试降低学习率或者加入梯度裁剪clip_grad_norm_。如果损失下降得非常慢可能是学习率太小如果震荡剧烈可能是学习率太大或批处理大小太小。模型输出范围确保你的模型最后一层使用了合适的激活函数。对于图像复原输出像素值应在[0,1]或[0,255]范围内。如果使用Tanh输出[-1,1]而你的目标图像是[0,1]损失函数计算就会出错。通常最后一层不使用激活函数线性层或者使用Sigmoid来约束到[0,1]。预训练权重加载如果你在微调预训练模型确认加载权重的键名是否匹配。使用strictFalse可以避免因结构微调导致的错误但务必打印出缺失和多余的键检查是否是你期望的。过拟合如果训练损失持续下降但验证损失很早就开始上升这是典型的过拟合。对策包括增加数据增强的强度、使用更小的模型、添加权重衰减weight_decay、使用Dropout如果模型支持、或者尽早停止训练Early Stopping。5.2 显存不足OOM的解决方案训练SwinIR时显存不足是家常便饭。除了换更大的显卡可以尝试以下方法减小批处理大小这是最直接有效的方法。但批处理大小太小可能会影响BatchNorm的统计稳定性并降低训练速度。减小输入图像尺寸在数据加载时使用更小的patch_size进行裁剪。例如从128x128降到64x64显存占用大致会降为1/4。梯度累积这是一种模拟大batch_size的技术。你以较小的batch_size进行前向传播和反向传播但不立即更新权重而是累积多个小批次的梯度当累积步数达到设定值后再进行一次更新。这不会减少单次前向传播的显存占用但允许你使用更小的物理batch_size。accumulation_steps 4 optimizer.zero_grad() for i, (data, target) in enumerate(dataloader): output model(data) loss criterion(output, target) / accumulation_steps # 损失按累积步数平均 loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()混合精度训练使用Automatic Mixed Precision (AMP)。这通过将部分计算转换为16位浮点数FP16来减少显存占用并加速计算。PyTorch中实现非常简单from torch.cuda.amp import autocast, GradScaler scaler GradScaler() ... with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意AMP有时可能导致数值不稳定如果出现NaN损失可能需要调整GradScaler的参数或关闭AMP。5.3 提升模型效果的实用技巧更优的数据增强除了基础的翻转旋转可以尝试更复杂的增强如颜色抖动Color Jitter、模糊Blur、添加噪声等这能提升模型的鲁棒性。但要注意增强必须同步应用于LR和HR图像对。学习率热身Warmup在训练开始时从一个很小的学习率线性增加到预设的初始学习率有助于稳定训练初期。这对于Transformer模型尤其有益。可以结合余弦退火调度器一起使用。指数移动平均EMA维护模型权重的一个滑动平均版本在验证和测试时使用这个平均版本通常能获得更稳定、泛化能力更强的模型。多尺度训练在训练时不是固定裁剪patch_size而是在一个范围内随机选择例如[64, 128, 192]这可以让模型学习到不同尺度的特征提升泛化能力。实现时需要在SRDataset的__getitem__中动态生成裁剪尺寸。损失函数组合如前所述结合L1损失、感知损失甚至对抗损失是提升视觉质量的必经之路。可以从L1感知损失开始权重系数需要仔细调整如L1:1.0, 感知:0.01。5.4 关于“全参训练与微调对显存要求的区别”的理解在相关热词中提到了这个问题这里简单解释一下。全参训练是指随机初始化模型所有权重从头开始训练。此时优化器需要为每一个参数保存一份状态例如Adam优化器需要保存一阶矩估计和二阶矩估计这需要额外的显存。微调通常指加载预训练模型的大部分权重可能冻结一部分层只训练另一部分如最后的重建头。此时被冻结层的参数不需要计算梯度优化器也不需要为其保存状态因此可以节省大量显存。在我们的代码框架中通过设置param.requires_grad False来冻结参数就能实现这一点。所以在相同模型和批处理大小下微调通常比全参训练占用更少的显存也收敛得更快。本文还有配套的精品资源点击获取 SwinIR图像复原模型自定义训练与测试代码实战指南
简介本资源是一套基于SwinIR模型的自定义图像恢复训练与测试代码实现面向计算机视觉方向的研究者与深度学习初学者聚焦图像超分辨率重建与去噪等典型低级视觉任务。代码逻辑完整、注释清晰支持开箱即用图像去噪任务仅需调整数据路径即可运行超分任务需在数据加载类中取消patchsize相关操作便于理解底层数据处理机制。压缩包共17个文件含5个核心Python源码如main.py、net.py、data.py、4个编译缓存pyc、4个IDE配置XML文件及MATLAB评估脚本psnr.m、compute_psnr.m等整体仅28KB轻量易读。已有5370人学习下载代码结构层次分明涵盖数据集组织inputTrain/targetTrain等目录、网络构建、损失定义、训练循环与结果评估全流程特别适合复现SwinTransformer在图像恢复中的应用并开展二次开发。1. 项目概述从“拿来就用”到“深度定制”的SwinIR训练之旅如果你在图像超分辨率、去噪或者JPEG压缩伪影去除这些任务上折腾过大概率听说过SwinIR这个名字。它基于Swin Transformer架构在多个图像复原的公开基准测试上都刷出了非常亮眼的成绩。很多朋友拿到官方代码仓库跑通预训练模型在标准数据集上的测试感觉效果不错但一到想用自己的数据、调整网络结构或者尝试新的损失函数时就卡住了。官方代码库为了通用性和可复现性往往结构复杂各种配置文件和抽象层叠在一起对于想快速验证一个想法或者进行教学演示的人来说学习成本太高。这就是我动手整理这份“逻辑完整易懂”的自定义训练测试代码的初衷——剥离所有不必要的抽象把数据流、模型前向传播、损失计算、优化器更新、验证测试这几个核心环节用最直白、模块化的Python代码呈现出来。让你不仅能跑起来更能看清楚每一行代码在做什么知道在哪里插入你自己的修改。这份代码的核心价值在于“透明”和“可插拔”。它不追求大而全的工程化而是聚焦于SwinIR模型训练与测试最本质的流程。你将看到数据是如何被加载、预处理并送入模型的损失函数是如何计算并反向传播的模型权重是如何保存和加载的。更重要的是代码结构清晰每个功能块数据加载、模型定义、训练循环、测试评估都是独立的你可以像搭积木一样替换其中的任何部分。例如想把L1损失换成感知损失Perceptual Loss只需要修改训练循环中的几行代码。想尝试不同的学习率调度策略替换对应的模块即可。这对于研究者快速进行算法迭代或者初学者深入理解深度学习训练流程都是非常有帮助的。2. 核心模块拆解你的数据、模型与训练循环一份完整的训练代码可以抽象为三个核心部分数据管道、模型定义和训练引擎。下面我们逐一拆解并解释在SwinIR这个具体场景下每一部分是如何实现的。2.1 数据加载器打造专属的图像对流水线图像复原任务如超分的训练数据通常是“低质量-高质量”图像对。我们的数据加载器Dataset和DataLoader需要高效地读取、配对并预处理这些图像。2.1.1 自定义Dataset类的设计逻辑我们首先定义一个SRDataset类。它的核心工作是给定一个包含低分辨率LR图像和高分辨率HR图像的文件夹路径在初始化时建立图像对的路径列表。这里的关键设计点是如何组织你的数据。一个常见的简单结构是your_dataset/ ├── LR/ # 存放所有低分辨率图像 │ ├── 001.png │ ├── 002.png │ └── ... └── HR/ # 存放所有对应的高分辨率图像 ├── 001.png ├── 002.png └── ...要求LR和HR文件夹中的文件名严格一一对应。SRDataset的__getitem__方法会接收一个索引然后同时读取对应索引的LR和HR图像。2.1.2 图像预处理与增强的关键步骤读取图像通常使用PIL或OpenCV后不能直接丢给模型。需要经过一系列预处理变换这些变换通过PyTorch的torchvision.transforms组合成一个transform管道。对于训练阶段典型的流程包括随机裁剪为了增加数据的多样性和实现批训练我们从HR图像中随机裁剪一个固定大小的块如patch_size128然后在LR图像对应的位置考虑缩放因子例如超分4倍则LR中裁剪128/432的块进行同步裁剪。这确保了图像对的区域一致性。随机旋转与翻转这是最常用且几乎无成本的几何数据增强。以一定概率对图像对进行90°、180°、270°的随机旋转以及水平/垂直翻转可以极大增加数据多样性防止模型过拟合。转换为张量将PIL图像或NumPy数组转换为PyTorch张量并将像素值范围从[0, 255]归一化到[0.0, 1.0]。可选通道归一化有时会进一步根据数据集的均值和标准差进行归一化但对于图像复原使用[0,1]范围通常已足够。对于测试/验证阶段变换则简单得多通常只需要将图像缩放到统一大小或保持原尺寸然后转换为张量和归一化。一个重要的细节测试时如果图像尺寸不能被模型的缩放因子整除可能会在边界产生伪影。因此有时需要对测试图像进行填充padding以满足整除条件在输出后再将填充部分裁剪掉。2.1.3 DataLoader的参数调优经验将定义好的Dataset交给DataLoader。有几个参数直接影响训练效率batch_size根据你的GPU显存决定。SwinIR模型不算小在消费级显卡如RTX 4090上超分任务可能batch_size只能设为8或16。可以从一个小值开始尝试逐步增加直到显存占满。num_workers数据加载的并行进程数。建议设置为CPU核心数或略少可以显著加速数据加载避免训练循环等待数据。在Linux系统上效果尤为明显。pin_memoryTrue当使用GPU时将此参数设为True可以将数据更快地从主机内存传输到GPU显存提升训练速度。shuffleTrue仅在训练时使用打乱数据顺序使每个epoch看到的数据分布都不同。注意如果你的数据集非常大例如数万对图像且存储于机械硬盘过高的num_workers可能会导致磁盘I/O成为瓶颈。此时需要平衡。使用SSD硬盘可以极大缓解此问题。2.2 模型定义理解SwinIR的骨干与头SwinIR的官方实现通常包含多个部分浅层特征提取、深层特征提取Swin Transformer块和高质量图像重建。在我们的简化版代码中我们会聚焦于核心部分。2.2.1 浅层特征提取从像素到特征这部分通常是一个简单的卷积层nn.Conv2d。它的作用是将输入的3通道RGB低质量图像映射到一个更高维的特征空间例如64或128个通道。你可以把它理解为一个“编码器”的入口它将原始的像素信息转换为一系列更抽象、更适合后续深度网络处理的特征图。这个卷积层的核大小通常是3x3步长为1填充为1以保持空间分辨率不变。2.2.2 核心骨干Swin Transformer块堆叠这是SwinIR的灵魂。浅层特征会被送入一系列Swin Transformer块Swin Transformer Blocks中。每个块内部主要包含窗口多头自注意力W-MSA将特征图划分为不重叠的局部窗口如8x8在每个窗口内计算自注意力。这大大降低了传统全局自注意力的计算复杂度从图像尺寸的平方级降低到窗口尺寸的平方级。移位窗口多头自注意力SW-MSA为了引入窗口间的信息交互在连续的Swin Transformer块中会交替使用W-MSA和SW-MSA。SW-MSA会将窗口进行一定偏移例如向右下角偏移半个窗口使得当前窗口能与上一层的不同窗口区域进行交互。多层感知机MLP在每个注意力层之后会接一个两层的全连接网络中间带有GELU激活函数用于进一步的特征变换。层归一化LayerNorm和残差连接每个子层注意力、MLP前都应用层归一化并广泛使用残差连接来缓解深层网络的梯度消失问题。在代码中这部分通常被组织为一个nn.ModuleList里面包含多个RSTBResidual Swin Transformer Block或类似命名的模块。一个关键的超参数是“深度”即堆叠多少个这样的块。更多的块意味着更大的模型容量和更长的训练时间但也可能带来性能提升需要根据任务和数据集权衡。2.2.3 重建模块从特征回到图像经过深层特征提取后我们得到了富含上下文信息的特征图。重建模块的任务是将这些特征图“翻译”回高质量的RGB图像。对于超分任务这通常是一个亚像素卷积层nn.PixelShuffle的上采样方式或一个转置卷积层nn.ConvTranspose2d后接一个最终的卷积层将通道数调整回3。对于去噪或去块效应任务可能只需要一个卷积层来直接预测残差即“干净图像噪声图像预测的噪声”或重建图像。2.2.4 预训练权重的加载与部分微调SwinIR官方提供了在多个大型数据集如DIV2K上预训练的模型。在我们的代码中加载预训练权重是一个非常重要的步骤可以让你在小数据集上快速获得好效果。使用model.load_state_dict(torch.load(pretrain_path), strictFalse)来加载。strictFalse参数很关键因为它允许只加载模型和预训练权重中名称匹配的部分对于你修改过的层例如改变了输出通道数可以忽略从而支持部分微调。如果你想冻结骨干网络只训练重建头可以这样做for name, param in model.named_parameters(): if ‘head’ not in name: # 假设重建层的参数名包含‘head’ param.requires_grad False这样在反向传播时只有requires_gradTrue的参数才会更新可以节省显存并加速训练初期收敛。2.3 训练循环损失、优化与评估的舞蹈这是整个代码的“发动机”一个标准的PyTorch训练循环但其中每一步都有讲究。2.3.1 损失函数的选择L1还是L2图像复原任务最常用的损失函数是L1损失MAE和L2损失MSE。L2损失MSE惩罚大的误差更重倾向于产生更平滑的输出但有时会导致图像过于模糊丢失高频纹理细节。L1损失MAE对误差的惩罚是线性的通常能产生更清晰的边缘和纹理是目前超分等任务的主流选择。在我们的代码中你可以简单地使用nn.L1Loss()。但更优的做法是结合多种损失。例如感知损失Perceptual Loss利用预训练的分类网络如VGG提取特征在特征空间计算差异使复原图像在“感知”上更接近真实图像提升视觉质量。对抗损失GAN Loss引入一个判别器网络让生成器SwinIR产生尽可能“以假乱真”的图像可以极大地提升纹理的真实感。但这会使得训练变得不稳定需要仔细调整。在简化版代码中我们可能只使用L1损失作为起点。你可以通过以下方式轻松扩展criterion_pixel nn.L1Loss() criterion_perceptual PerceptualLoss() # 需要自定义或引用第三方库 ... loss criterion_pixel(output, target) 0.01 * criterion_perceptual(output, target)2.3.2 优化器与学习率调度寻找下降的路径对于SwinIR这类Transformer-based模型Adam或AdamW优化器是标准选择。AdamW相比Adam加入了权重衰减的正则化通常泛化性能更好。optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4)初始学习率lr是一个关键超参数。对于微调预训练模型通常使用较小的学习率如1e-5到1e-4对于从头训练可以稍大如1e-4。学习率调度同样重要。最常见的是余弦退火调度CosineAnnealingLR它让学习率随着训练过程从初始值平滑地衰减到0模拟了“精细调参”的过程。另一种实用的策略是带热重启的余弦退火CosineAnnealingWarmRestarts它周期性地重启学习率有助于模型跳出局部最优。在我们的代码框架中你可以很方便地集成这些调度器。2.3.3 训练迭代中的关键操作一个训练迭代iteration包含以下步骤清零梯度optimizer.zero_grad()。这是必须的因为PyTorch会累积梯度。前向传播outputs model(lr_images)。计算损失loss criterion(outputs, hr_images)。反向传播loss.backward()。计算所有可训练参数关于损失的梯度。梯度裁剪可选但推荐torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。这对于稳定Transformer模型的训练尤其重要可以防止梯度爆炸。参数更新optimizer.step()。根据梯度和优化算法更新模型权重。学习率更新scheduler.step()在每个epoch或iteration后取决于调度器类型。2.3.4 验证与模型保存避免过拟合的眼睛我们通常在每个epoch结束后在独立的验证集上评估模型性能。验证时必须将模型设置为评估模式model.eval()这会关闭Dropout、BatchNorm等层在训练和评估时的不同行为。同时要使用torch.no_grad()上下文管理器来禁用梯度计算节省内存和计算资源。评估指标除了损失值还应包括感知指标例如PSNR峰值信噪比最常用的客观指标值越高越好但与主观视觉质量有时不完全一致。SSIM结构相似性比PSNR更能反映人眼对结构信息的感知。模型保存策略通常有两种定期保存每N个epoch保存一次用于检查点恢复。保存最佳模型根据验证集上的PSNR或SSIM只保存指标最好的那个模型权重。这是最常用的策略。if current_psnr best_psnr: best_psnr current_psnr torch.save(model.state_dict(), ‘best_model.pth’)3. 测试与推理让模型真正工作起来训练好的模型最终要用于处理新的图像。测试脚本需要独立、高效且易于使用。3.1 单图与批处理推理脚本推理脚本的核心流程是加载模型权重 - 预处理输入图像 - 模型前向传播 - 后处理输出图像 - 保存。3.1.1 图像预处理对齐训练时配置这一点至关重要测试时对输入图像的预处理裁剪、归一化等必须与训练时完全一致。如果训练时输入是[0,1]范围测试时读入的[0,255]图像就必须除以255。如果训练时对数据进行了减均值除标准差的操作测试时也必须用相同的均值和标准差。不一致的预处理会导致模型性能严重下降甚至产生毫无意义的结果。3.1.2 处理任意尺寸的输入训练时我们使用固定大小的裁剪块但测试时图像尺寸千变万化。SwinIR模型由于其中的Patch Merging/Embedding等操作可能对输入尺寸有整除要求例如需要是窗口大小的整数倍。一个常见的技巧是反射填充import torch.nn.functional as F mod_pad_h (window_size - h % window_size) % window_size mod_pad_w (window_size - w % window_size) % window_size img F.pad(img, (0, mod_pad_w, 0, mod_pad_h), ‘reflect’)在模型输出后再将填充的部分裁剪掉得到最终结果。3.1.3 批处理与GPU内存管理当需要处理大量图像或视频序列时批处理可以提升效率。但要注意大尺寸图像进行批处理会消耗巨量显存。一个稳健的策略是动态调整批处理大小。可以先尝试一个较小的batch_size如1或2如果GPU内存允许再逐步增加。也可以编写一个循环每次处理一批并适时清空缓存 (torch.cuda.empty_cache())。3.2 性能评估与可视化对比单纯的PSNR/SSIM数字有时不够直观。一个好的测试脚本应该能生成可视化的对比图。3.2.1 生成对比网格图将输入的低质量图像、模型输出的复原图像、真实的高质量图像如果有的话并排放在一起保存为一张图片。这能最直观地展示模型的效果。可以使用Matplotlib或OpenCV来轻松实现。更进一步可以计算并标注出每张结果图与GT之间的PSNR/SSIM值。3.2.2 局部区域放大对比人眼对纹理和边缘细节非常敏感。全局图看起来可能都不错但放大到某个特定区域如文字、毛发、建筑纹理差异就显现了。在生成对比图时可以额外插入一个放大特定区域的子图让优劣一目了然。3.2.3 指标批量计算与报告对于整个测试集脚本应该能自动遍历所有图像计算平均PSNR、平均SSIM并生成一个简洁的文本报告或CSV文件。这对于量化比较不同模型或不同训练轮次的效果至关重要。4. 代码实战从零搭建可运行的训练流程让我们抛开复杂的抽象用最直接的代码勾勒出骨架。这里会省略一些非常具体的Swin Transformer块实现细节你可以从官方仓库引用但会完整展示所有接口和流程。4.1 数据准备与Dataset类实现首先确保你的数据按前述结构整理好。然后实现SRDatasetimport os from PIL import Image import torch from torch.utils.data import Dataset import torchvision.transforms as transforms class SRDataset(Dataset): def __init__(self, lr_dir, hr_dir, patch_size128, scale4, is_trainTrue): self.lr_dir lr_dir self.hr_dir hr_dir self.patch_size patch_size self.scale scale self.is_train is_train # 获取所有配对的文件名假设文件名相同 self.lr_filenames sorted([f for f in os.listdir(lr_dir) if f.endswith(‘.png’)]) self.hr_filenames sorted([f for f in os.listdir(hr_dir) if f.endswith(‘.png’)]) # 简单检查 assert len(self.lr_filenames) len(self.hr_filenames) for lr, hr in zip(self.lr_filenames, self.hr_filenames): assert lr hr, f“文件名不匹配: {lr} vs {hr}” # 定义变换 if self.is_train: self.transform transforms.Compose([ transforms.RandomCrop(patch_size), # 对HR进行随机裁剪 transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(90), transforms.ToTensor(), # 自动转换到[0,1] ]) else: self.transform transforms.Compose([ transforms.ToTensor(), ]) def __len__(self): return len(self.lr_filenames) def __getitem__(self, idx): lr_path os.path.join(self.lr_dir, self.lr_filenames[idx]) hr_path os.path.join(self.hr_dir, self.hr_filenames[idx]) lr_img Image.open(lr_path).convert(‘RGB’) hr_img Image.open(hr_path).convert(‘RGB’) # 注意对于训练我们需要同步裁剪LR和HR if self.is_train: # 对HR进行随机裁剪 i, j, h, w transforms.RandomCrop.get_params(hr_img, output_size(self.patch_size, self.patch_size)) hr_img_crop transforms.functional.crop(hr_img, i, j, h, w) # 对LR在对应位置进行裁剪考虑缩放因子 lr_img_crop transforms.functional.crop(lr_img, i//self.scale, j//self.scale, h//self.scale, w//self.scale) # 应用其他增强翻转、旋转 seed torch.randint(0, 2**32, (1,)).item() torch.manual_seed(seed) lr_img_crop self.transform(lr_img_crop) torch.manual_seed(seed) # 确保LR和HR应用完全相同的随机变换 hr_img_crop self.transform(hr_img_crop) return lr_img_crop, hr_img_crop else: lr_img self.transform(lr_img) hr_img self.transform(hr_img) return lr_img, hr_img, self.lr_filenames[idx] # 测试时返回文件名用于保存结果4.2 简化版SwinIR模型定义与训练循环这里我们定义一个极简的模型外壳重点展示训练循环的逻辑。假设我们已经有一个叫SwinIR的模型类可以从官方实现中导入或简化实现。import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm # 用于显示进度条 import time # 假设我们已经定义或导入了SwinIR模型 # from models.network_swinir import SwinIR class SimpleSwinIR(nn.Module): # 此处应为具体的模型定义这里用占位符 def __init__(self, upscale4): super().__init__() # 这里应包含浅层特征提取、Swin Transformer块、重建头等 self.conv_first nn.Conv2d(3, 64, 3, 1, 1) # ... 更多层定义 self.recon_head nn.Conv2d(64, 3, 3, 1, 1) def forward(self, x): # ... 前向传播逻辑 return self.recon_head(x) def train_one_epoch(model, dataloader, criterion, optimizer, device, epoch): model.train() running_loss 0.0 progress_bar tqdm(dataloader, descf‘Epoch [{epoch}]‘) for batch_idx, (lr_imgs, hr_imgs) in enumerate(progress_bar): lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) # 清零梯度 optimizer.zero_grad() # 前向传播 outputs model(lr_imgs) # 计算损失 loss criterion(outputs, hr_imgs) # 反向传播 loss.backward() # 梯度裁剪可选但推荐 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 更新参数 optimizer.step() running_loss loss.item() progress_bar.set_postfix({‘loss’: loss.item()}) avg_loss running_loss / len(dataloader) return avg_loss def validate(model, dataloader, criterion, device): model.eval() running_loss 0.0 total_psnr 0.0 with torch.no_grad(): for lr_imgs, hr_imgs in dataloader: lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) outputs model(lr_imgs) loss criterion(outputs, hr_imgs) running_loss loss.item() # 计算PSNR (假设图像范围在[0,1]) mse torch.mean((outputs - hr_imgs) ** 2) psnr 20 * torch.log10(1.0 / torch.sqrt(mse)) total_psnr psnr.item() avg_loss running_loss / len(dataloader) avg_psnr total_psnr / len(dataloader) return avg_loss, avg_psnr def main(): # 配置参数 device torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’) lr_dir ‘path/to/your/LR/train’ hr_dir ‘path/to/your/HR/train’ val_lr_dir ‘path/to/your/LR/val’ val_hr_dir ‘path/to/your/HR/val’ batch_size 16 num_epochs 100 learning_rate 1e-4 # 数据加载 train_dataset SRDataset(lr_dir, hr_dir, patch_size128, scale4, is_trainTrue) val_dataset SRDataset(val_lr_dir, val_hr_dir, patch_size128, scale4, is_trainFalse) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size1, shuffleFalse, num_workers2, pin_memoryTrue) # 模型、损失、优化器 model SimpleSwinIR(upscale4).to(device) criterion nn.L1Loss() optimizer optim.AdamW(model.parameters(), lrlearning_rate, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxnum_epochs) best_psnr 0.0 for epoch in range(1, num_epochs1): # 训练一个epoch train_loss train_one_epoch(model, train_loader, criterion, optimizer, device, epoch) # 验证 val_loss, val_psnr validate(model, val_loader, criterion, device) print(f‘Epoch {epoch}: Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val PSNR: {val_psnr:.2f}dB’) # 学习率调度 scheduler.step() # 保存最佳模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(model.state_dict(), f‘best_model_epoch{epoch}_psnr{val_psnr:.2f}.pth’) print(f‘模型已保存 (PSNR: {val_psnr:.2f}dB)’) # 定期保存检查点 if epoch % 10 0: torch.save({ ‘epoch’: epoch, ‘model_state_dict’: model.state_dict(), ‘optimizer_state_dict’: optimizer.state_dict(), ‘scheduler_state_dict’: scheduler.state_dict(), ‘best_psnr’: best_psnr, }, f‘checkpoint_epoch{epoch}.pth’) if __name__ ‘__main__’: main()4.3 测试脚本示例训练完成后使用以下脚本进行单张图像测试import torch from PIL import Image import torchvision.transforms as transforms import numpy as np import cv2 def load_model(model_path, device): model SimpleSwinIR(upscale4).to(device) state_dict torch.load(model_path, map_locationdevice) model.load_state_dict(state_dict) model.eval() return model def preprocess_image(image_path, scale4, window_size8): “”“预处理单张图像包括填充以满足窗口大小整除要求”“” img Image.open(image_path).convert(‘RGB’) img_tensor transforms.ToTensor()(img).unsqueeze(0) # [1, C, H, W] # 填充 _, _, h, w img_tensor.shape mod_pad_h (window_size - h % window_size) % window_size mod_pad_w (window_size - w % window_size) % window_size img_tensor torch.nn.functional.pad(img_tensor, (0, mod_pad_w, 0, mod_pad_h), ‘reflect’) return img_tensor, h, w def test_single_image(model, lr_image_path, output_path, device, scale4): model.eval() with torch.no_grad(): input_tensor, orig_h, orig_w preprocess_image(lr_image_path, scale) input_tensor input_tensor.to(device) output_tensor model(input_tensor) # 裁剪掉填充的部分 output_tensor output_tensor[:, :, :orig_h*scale, :orig_w*scale] # 后处理张量转回图像 output_np output_tensor.squeeze(0).cpu().numpy() # [C, H, W] output_np np.transpose(output_np, (1, 2, 0)) # [H, W, C] output_np np.clip(output_np * 255.0, 0, 255).astype(np.uint8) # 保存 cv2.imwrite(output_path, cv2.cvtColor(output_np, cv2.COLOR_RGB2BGR)) print(f‘结果已保存至: {output_path}’) # 使用示例 device torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’) model load_model(‘best_model.pth’, device) test_single_image(model, ‘test_input.png’, ‘test_output.png’, device)5. 避坑指南与性能调优实战心得在实际操作中你会遇到各种各样的问题。下面是我在多次训练SwinIR和类似模型时总结的一些关键经验和常见陷阱。5.1 训练不收敛或效果差的排查清单如果你的模型训练后效果很差甚至不收敛请按以下顺序检查数据与预处理这是最常见的问题源。首先确保你的LR-HR图像对是严格对齐的。用眼睛看几对放大检查边缘是否对齐。其次检查预处理代码确保训练和测试时的归一化方式是[0,1]还是[0,255]是否减了均值完全一致。一个快速验证方法是取一张训练图像经过Dataset的__getitem__处理后再保存出来看看是否还是正常的图像。损失函数值观察训练初期的损失值。如果损失从一开始就是NaN或无限大很可能是梯度爆炸。尝试降低学习率或者加入梯度裁剪clip_grad_norm_。如果损失下降得非常慢可能是学习率太小如果震荡剧烈可能是学习率太大或批处理大小太小。模型输出范围确保你的模型最后一层使用了合适的激活函数。对于图像复原输出像素值应在[0,1]或[0,255]范围内。如果使用Tanh输出[-1,1]而你的目标图像是[0,1]损失函数计算就会出错。通常最后一层不使用激活函数线性层或者使用Sigmoid来约束到[0,1]。预训练权重加载如果你在微调预训练模型确认加载权重的键名是否匹配。使用strictFalse可以避免因结构微调导致的错误但务必打印出缺失和多余的键检查是否是你期望的。过拟合如果训练损失持续下降但验证损失很早就开始上升这是典型的过拟合。对策包括增加数据增强的强度、使用更小的模型、添加权重衰减weight_decay、使用Dropout如果模型支持、或者尽早停止训练Early Stopping。5.2 显存不足OOM的解决方案训练SwinIR时显存不足是家常便饭。除了换更大的显卡可以尝试以下方法减小批处理大小这是最直接有效的方法。但批处理大小太小可能会影响BatchNorm的统计稳定性并降低训练速度。减小输入图像尺寸在数据加载时使用更小的patch_size进行裁剪。例如从128x128降到64x64显存占用大致会降为1/4。梯度累积这是一种模拟大batch_size的技术。你以较小的batch_size进行前向传播和反向传播但不立即更新权重而是累积多个小批次的梯度当累积步数达到设定值后再进行一次更新。这不会减少单次前向传播的显存占用但允许你使用更小的物理batch_size。accumulation_steps 4 optimizer.zero_grad() for i, (data, target) in enumerate(dataloader): output model(data) loss criterion(output, target) / accumulation_steps # 损失按累积步数平均 loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()混合精度训练使用Automatic Mixed Precision (AMP)。这通过将部分计算转换为16位浮点数FP16来减少显存占用并加速计算。PyTorch中实现非常简单from torch.cuda.amp import autocast, GradScaler scaler GradScaler() ... with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意AMP有时可能导致数值不稳定如果出现NaN损失可能需要调整GradScaler的参数或关闭AMP。5.3 提升模型效果的实用技巧更优的数据增强除了基础的翻转旋转可以尝试更复杂的增强如颜色抖动Color Jitter、模糊Blur、添加噪声等这能提升模型的鲁棒性。但要注意增强必须同步应用于LR和HR图像对。学习率热身Warmup在训练开始时从一个很小的学习率线性增加到预设的初始学习率有助于稳定训练初期。这对于Transformer模型尤其有益。可以结合余弦退火调度器一起使用。指数移动平均EMA维护模型权重的一个滑动平均版本在验证和测试时使用这个平均版本通常能获得更稳定、泛化能力更强的模型。多尺度训练在训练时不是固定裁剪patch_size而是在一个范围内随机选择例如[64, 128, 192]这可以让模型学习到不同尺度的特征提升泛化能力。实现时需要在SRDataset的__getitem__中动态生成裁剪尺寸。损失函数组合如前所述结合L1损失、感知损失甚至对抗损失是提升视觉质量的必经之路。可以从L1感知损失开始权重系数需要仔细调整如L1:1.0, 感知:0.01。5.4 关于“全参训练与微调对显存要求的区别”的理解在相关热词中提到了这个问题这里简单解释一下。全参训练是指随机初始化模型所有权重从头开始训练。此时优化器需要为每一个参数保存一份状态例如Adam优化器需要保存一阶矩估计和二阶矩估计这需要额外的显存。微调通常指加载预训练模型的大部分权重可能冻结一部分层只训练另一部分如最后的重建头。此时被冻结层的参数不需要计算梯度优化器也不需要为其保存状态因此可以节省大量显存。在我们的代码框架中通过设置param.requires_grad False来冻结参数就能实现这一点。所以在相同模型和批处理大小下微调通常比全参训练占用更少的显存也收敛得更快。本文还有配套的精品资源点击获取