尧图网页设计 
简介本资源是图像超分辨率领域经典模型SRGAN的PyTorch完整复现工程面向计算机视觉方向的初学者与进阶研究者聚焦真实场景下的单图超分任务尤其适用于课程设计、科研复现实验及轻量级项目集成。压缩包共375个文件含10个核心Python脚本如train.py、test_image.py、draw_evaluation.py等、324张测试/验证用PNG图像涵盖barbara、lenna、baboon等标准基准图、3个最优PSNR模型权重x2/x4/x8倍率及训练过程产出的epochs、statistics、training_results等结构化结果目录整体大小231.09MB。已有694人学习下载资源配套CSDN技术博文详解实现逻辑与调参要点。用户可直接运行demo.py进行GT/Bicubic/SRGAN三栏可视化对比调用test_video.py处理视频序列并通过draw_evaluation.py一键生成Loss、PSNR、SSIM随Epoch变化曲线所有模块注释详尽、路径规范、即插即用。1. 项目概述从SRGAN到可复现的PyTorch工程如果你在图像处理领域摸爬滚打过一阵子肯定对“超分辨率”这个词不陌生。简单说就是把一张模糊的、像素低的小图通过算法“脑补”出细节变成一张清晰的高清大图。这技术从早期的插值算法到后来的深度学习模型发展速度飞快。而SRGAN可以说是这个领域里一个里程碑式的存在。它第一次把生成对抗网络GAN的思路引入超分辨率不再仅仅追求像素值上的接近比如PSNR指标而是开始关注人眼视觉感受生成纹理更真实、细节更丰富的图像。我最初接触SRGAN时是在读它的论文。理论很优美但真到了自己动手复现的时候才发现从论文到代码之间隔着一片“坑海”。官方没有提供完整的训练代码网上的开源实现质量参差不齐有的注释不清有的依赖老旧有的甚至关键步骤都有错误。训练过程更是玄学GAN本身就不稳定再加上超分辨率任务对细节的极致追求调参调得人头皮发麻。更别提还要自己写训练日志、画损失曲线、保存最优模型这些工程化的事情了一套流程下来精力全耗在搭轮子上反而没多少时间去理解模型本身了。所以当我决定重新梳理并实现一个SRGAN的PyTorch版本时目标就非常明确不仅要代码能跑通更要让它成为一个“教学级”和“生产级”兼备的项目。这意味着每一行核心代码都要有详尽的注释解释清楚“为什么这么做”训练过程要可视化让你能实时看到模型是学好了还是在“摆烂”还要自动保存PSNR指标最好的模型权重方便你直接拿来用或者继续微调。最终这个项目覆盖了x2、x4、x8三种常见的超分倍数你拿到手就能在自己的数据上跑起来或者以此为蓝本去实现更复杂的变体。2. 核心思路与方案选型为什么是PyTorch SRGAN2.1 模型架构的再审视生成器与判别器的博弈SRGAN的核心思想是“对抗训练”。它包含两个网络生成器Generator它的任务是把低分辨率LR图像“想象”成高分辨率HR图像。在SRGAN中生成器的主体是一个深度残差网络。为什么用残差结构因为超分任务可以看作是在LR图像基础上添加高频细节残差残差连接能有效缓解深层网络的梯度消失问题让网络专注于学习LR到HR的“差异部分”训练起来更稳定、更快。判别器Discriminator它是一个二分类器任务就是判断一张图像是“真实的”高清图来自数据集还是“伪造的”由生成器产生的高清图。它的存在就像一位严厉的考官不断逼迫生成器去生成以假乱真的图像。两者的关系是动态博弈生成器努力骗过判别器判别器努力识破生成器。这个博弈过程最终驱使生成器产生的图像不仅在像素值上接近真实在纹理、边缘等视觉感知上也更加逼真。这与传统只使用MSE均方误差损失的方法有本质区别。MSE损失容易导致结果过于平滑丢失纹理因为它是“平均主义者”而GAN的对抗损失鼓励生成“局部逼真”的细节。注意原始SRGAN论文中生成器的损失函数是内容损失Content Loss、**对抗损失Adversarial Loss和感知损失Perceptual Loss基于VGG特征**的加权和。其中感知损失是让生成图像和真实图像在深层特征空间而非像素空间上接近这是提升视觉质量的关键。我们的复现严格遵循了这一设计。2.2 为什么选择PyTorch作为实现框架面对TensorFlow、PyTorch乃至JAX等框架我选择了PyTorch这是基于多年一线开发经验的考量动态图优先PyTorch的 eager execution 模式让调试变得异常直观。你可以在任意地方插入print或使用调试器查看张量的形状和值这对于理解GAN这种复杂、动态的训练过程至关重要。想象一下当判别器损失突然变成NaN时你能快速定位到是梯度爆炸还是数据出了问题。生态与社区PyTorch在学术研究和工业界的普及度已毋庸置疑。这意味着你遇到的几乎所有问题都能在社区找到讨论和解决方案。从最新的模型实现如ESRGAN、Real-ESRGAN到各种数据加载、可视化工具PyTorch的生态是最丰富的。简洁的API设计torch.nn.Module和torch.optim的设计非常Pythonic定义模型和训练循环的逻辑清晰直白。自定义层、损失函数都轻而易举这让实现SRGAN中特定的残差块Residual Block和感知损失变得很顺畅。部署友好虽然本项目侧重于训练和复现但PyTorch模型通过TorchScript或ONNX可以相对方便地转换为后续移动端或服务端部署留下了可能性。相比之下虽然TensorFlow 2.x也改善了易用性但其历史包袱和某些API设计仍显繁琐。对于这样一个以清晰、教育为目的的复现项目PyTorch是更自然的选择。2.3 项目特色与解决的问题市面上有很多SRGAN代码那这个复现项目有什么不同它主要解决了以下几个痛点注释的深度注释不仅仅是解释“这行代码在做什么”What更重要的是解释“为什么要这么做”Why。例如在初始化权重时会说明为何使用kaiming_normal_初始化而不是xavier_uniform_在训练循环中会解释为何要先更新判别器再更新生成器。训练过程透明化训练GAN就像驾驶一架仪表盘不全的飞机。本项目集成了TensorBoard或Matplotlib实时绘图将生成器损失、判别器损失、PSNR、SSIM等关键指标的变化曲线实时可视化。你能一眼看出模型是否收敛是否发生了模式崩溃Mode Collapse。模型管理的自动化我们不止在每一个epoch结束时保存模型。代码会自动监控验证集上的PSNR指标只保存达到历史最优PSNR的那个模型权重。同时也会保存最新的模型作为checkpoint防止训练意外中断。这意味着你训练结束后直接就能拿到“表现最好”的模型而不是“最后一代”可能已经过拟合的模型。开箱即用的配置提供了x2, x4, x8三种上采样倍数的预训练模型或训练脚本配置。数据集处理、数据增强翻转、旋转、学习率调度MultiStepLR等常用技巧都已集成并配有详细的配置文件如YAML你只需要修改数据路径和几个参数就能跑起来。3. 环境搭建与核心代码解析3.1 从零开始PyTorch与CUDA环境配置工欲善其事必先利其器。一个稳定、版本匹配的深度学习环境是成功的第一步。很多人在这里踩坑问题多半出在CUDA、cuDNN和PyTorch的版本不匹配上。步骤一确认你的GPU和CUDA驱动打开命令行输入nvidia-smi。记下最上方显示的CUDA Version例如12.4。这是你的驱动支持的最高CUDA运行时版本但PyTorch需要安装与之兼容的CUDA Toolkit。步骤二使用Conda创建独立环境强烈建议使用Conda管理环境避免包冲突。conda create -n srgan_pytorch python3.9 conda activate srgan_pytorch选择Python 3.9是因为它在稳定性和对新旧包的兼容性上比较平衡。步骤三安装匹配的PyTorch这是最关键的一步。不要直接用pip install torch。前往 PyTorch官网 根据你的系统、Conda环境以及上一步查到的CUDA驱动版本选择对应的安装命令。 例如对于CUDA 12.1你可能会看到pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121对于没有NVIDIA GPU的机器则安装CPU版本。步骤四安装项目依赖在项目根目录下通常有一个requirements.txt文件。pip install -r requirements.txt典型的依赖包括opencv-python图像处理、pillow图像读取、tensorboard或tensorboardX可视化、matplotlib绘图、scikit-image计算PSNR/SSIM、tqdm进度条等。实操心得我习惯在安装完PyTorch后立刻写一个测试脚本验证GPU是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()返回False大概率是CUDA版本不匹配需要重新安装PyTorch或升级NVIDIA驱动。3.2 生成器网络结构深度拆解SRGAN的生成器是一个深层的残差网络。我们来逐层解析其PyTorch实现的关键部分。import torch.nn as nn import torch.nn.functional as F class ResidualBlock(nn.Module): 残差块两个卷积层 跳跃连接用于学习高频细节残差 def __init__(self, channels): super(ResidualBlock, self).__init__() self.conv1 nn.Conv2d(channels, channels, kernel_size3, padding1, biasFalse) self.bn1 nn.BatchNorm2d(channels) self.prelu nn.PReLU() # 使用PReLU而非ReLU允许小的负斜率可能提升性能 self.conv2 nn.Conv2d(channels, channels, kernel_size3, padding1, biasFalse) self.bn2 nn.BatchNorm2d(channels) def forward(self, x): residual x out self.conv1(x) out self.bn1(out) out self.prelu(out) out self.conv2(out) out self.bn2(out) # 残差连接将输入x加到输出上 out out residual return out class UpsampleBlock(nn.Module): 上采样块用于将特征图空间分辨率扩大2倍 def __init__(self, in_channels, up_scale): super(UpsampleBlock, self).__init__() # 使用PixelShuffle进行高效上采样 self.conv nn.Conv2d(in_channels, in_channels * (up_scale ** 2), kernel_size3, padding1) self.pixel_shuffle nn.PixelShuffle(up_scale) # 重排像素实现上采样 self.prelu nn.PReLU() def forward(self, x): out self.conv(x) out self.pixel_shuffle(out) out self.prelu(out) return out class Generator(nn.Module): SRGAN生成器网络 def __init__(self, scale_factor4, num_residual_blocks16): super(Generator, self).__init__() # 初始卷积层将输入图像映射到特征空间 self.conv1 nn.Sequential( nn.Conv2d(3, 64, kernel_size9, padding4), nn.PReLU() ) # 一系列残差块用于深度特征提取和学习残差 residual_blocks [] for _ in range(num_residual_blocks): residual_blocks.append(ResidualBlock(64)) self.residual_blocks nn.Sequential(*residual_blocks) # 残差块后的卷积层 self.conv2 nn.Sequential( nn.Conv2d(64, 64, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(64) ) # 上采样模块根据放大倍数决定上采样次数 upsample_blocks [] # 计算需要几次2倍上采样。例如 scale_factor4 - num_up2 num_up int(math.log2(scale_factor)) for _ in range(num_up): upsample_blocks.append(UpsampleBlock(64, up_scale2)) self.upsample_blocks nn.Sequential(*upsample_blocks) # 最终输出层将特征映射回RGB图像 self.conv3 nn.Conv2d(64, 3, kernel_size9, padding4) def forward(self, x): # 初始特征提取 out1 self.conv1(x) # 通过残差块 out self.residual_blocks(out1) # 残差连接大残差将初始特征与残差块输出相加 out self.conv2(out) out out out1 # 上采样 out self.upsample_blocks(out) # 最终输出使用Tanh将值约束到[-1, 1]与归一化后的输入数据匹配 out torch.tanh(self.conv3(out)) return out关键点解析PixelShuffle上采样这是SRGAN论文中的关键。传统上采样如转置卷积Deconv容易产生棋盘伪影Checkerboard Artifacts。PixelShuffle通过卷积增加通道数然后进行像素重排来扩大空间尺寸能生成更平滑的结果。例如将(C, H, W)的特征图通过卷积变为(C*4, H, W)再通过PixelShuffle(2)重排为(C, H*2, W*2)。残差连接的作用代码中有两处残差连接。一是在每个ResidualBlock内部保证梯度流动二是在所有残差块之后out out out1这是一个“大残差”连接让网络更容易学习到LR与HR之间的残差即细节信息而非学习完整的图像映射这大大降低了学习难度。输出激活函数Tanh输入图像在预处理时通常被归一化到[-1, 1]范围因此输出层使用Tanh将像素值约束到同一范围便于计算损失。3.3 判别器与损失函数的设计判别器是一个相对标准的CNN分类器但针对图像任务做了一些设计。class Discriminator(nn.Module): SRGAN判别器网络一个二分类CNN def __init__(self): super(Discriminator, self).__init__() self.net nn.Sequential( # 输入: [3, 96, 96] (假设HR patch大小) nn.Conv2d(3, 64, kernel_size3, padding1), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(64, 64, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(64), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸减半 nn.Conv2d(64, 128, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(128, 128, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸再减半 nn.Conv2d(128, 256, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(256, 256, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸再减半 nn.Conv2d(256, 512, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(512, 512, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplaceTrue), # 此时特征图已很小 nn.AdaptiveAvgPool2d(1), # 全局平均池化得到一个512维的向量 nn.Conv2d(512, 1024, kernel_size1), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(1024, 1, kernel_size1), # 输出一个标量通过Sigmoid表示“真”的概率 ) def forward(self, x): out self.net(x) return out.view(-1, 1) # 展平为 [batch_size, 1]损失函数的组合是SRGAN的灵魂。生成器的总损失是三项的加权和import torch import torch.nn as nn from torchvision.models import vgg19 class GeneratorLoss(nn.Module): def __init__(self, content_weight1e-2, adversarial_weight1e-3): super(GeneratorLoss, self).__init__() # 1. 内容损失 (MSE Loss) - 保证像素级相似性 self.mse_loss nn.MSELoss() # 2. 对抗损失 (BCE Loss) - 鼓励生成器骗过判别器 self.adversarial_loss nn.BCEWithLogitsLoss() # 注意判别器输出未经过Sigmoid故用WithLogits # 3. 感知损失 (Perceptual Loss) - 基于VGG特征保证高级语义相似性 vgg vgg19(pretrainedTrue).features[:36] # 取VGG19的前36层到conv5_4 for param in vgg.parameters(): param.requires_grad False # 冻结VGG参数 self.vgg vgg self.vgg_loss nn.MSELoss() self.content_weight content_weight self.adversarial_weight adversarial_weight def forward(self, sr_images, hr_images, discriminator_fake_output): # 计算MSE Loss mse_l self.mse_loss(sr_images, hr_images) # 计算对抗Loss目标是让判别器认为生成图像为“真”(label1) adversarial_l self.adversarial_loss(discriminator_fake_output, torch.ones_like(discriminator_fake_output)) # 计算感知Loss提取VGG特征 sr_features self.vgg(sr_images) hr_features self.vgg(hr_images) perceptual_l self.vgg_loss(sr_features, hr_features) # 加权总损失 total_loss mse_l self.content_weight * perceptual_l self.adversarial_weight * adversarial_l return total_loss, mse_l, perceptual_l, adversarial_l权重选择经验content_weight和adversarial_weight的平衡至关重要。对抗权重太大训练容易不稳定图像可能出现奇怪纹理太小则退化回普通MSE模型图像平滑。论文中分别使用1e-2和1e-3这是一个很好的起点。在实际训练中我通常会先以较大的内容损失权重如1.0和极小的对抗损失权重如1e-4预训练生成器几十个epoch让生成器先学会一个“基础版”的超分然后再调高对抗权重进行“对抗精炼”这样训练更稳定。4. 数据准备与训练流程实战4.1 数据集处理与高效数据加载高质量的数据是训练成功的一半。对于超分辨率任务你需要成对的低分辨率LR和高分辨率HR图像。通常有两种获取方式使用现成数据集如DIV2K、Flickr2K等它们已经提供了HR图像。你需要在线下用双三次下采样Bicubic Downsampling生成对应的LR图像。关键点下采样的插值算法必须与你的评估标准如PSNR所使用的算法一致通常是MATLAB的imresize函数。为了复现论文结果我强烈建议使用论文作者提供的脚本或opencv设置INTER_CUBIC来生成LR图并确保scale_factor准确。自定义数据如果你有自己的高清图库处理流程同上。数据加载器的构建from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as transforms import os class SRDataset(Dataset): def __init__(self, hr_dir, lr_dir, patch_size96, scale4, is_trainTrue): self.hr_dir hr_dir self.lr_dir lr_dir self.scale scale self.patch_size patch_size self.is_train is_train self.hr_images sorted([os.path.join(hr_dir, x) for x in os.listdir(hr_dir) if x.endswith((.png, .jpg))]) self.lr_images sorted([os.path.join(lr_dir, x) for x in os.listdir(lr_dir) if x.endswith((.png, .jpg))]) # 训练和验证/测试的数据转换不同 if is_train: self.transform transforms.Compose([ transforms.RandomCrop(patch_size), # 随机裁剪patch transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) # 归一化到[-1,1] ]) else: self.transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) ]) def __len__(self): return len(self.hr_images) def __getitem__(self, idx): hr_img Image.open(self.hr_images[idx]).convert(RGB) lr_img Image.open(self.lr_images[idx]).convert(RGB) # 确保LR和HR尺寸匹配 (HR LR * scale) lr_width, lr_height lr_img.size hr_width, hr_height hr_img.size assert hr_width lr_width * self.scale and hr_height lr_height * self.scale, \ f尺寸不匹配: LR {lr_img.size}, HR {hr_img.size}, scale {self.scale} 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_cropped transforms.functional.crop(hr_img, i, j, h, w) # 对LR进行对应区域的裁剪 (i//scale, j//scale, h//scale, w//scale) lr_img_cropped transforms.functional.crop(lr_img, i//self.scale, j//self.scale, h//self.scale, w//self.scale) hr_img hr_img_cropped lr_img lr_img_cropped hr_tensor self.transform(hr_img) lr_tensor self.transform(lr_img) return lr_tensor, hr_tensor重要细节Patch训练由于高清图像很大无法整张送入网络。通常随机裁剪出固定大小如96x96的HR patch并对应裁剪出LR patch进行训练。这增加了数据的多样性。数据增强随机水平/垂直翻转是基本操作有时还会加入随机旋转90, 180, 270度来进一步增强。注意增强操作应同时应用于配对的LR和HR patch保持空间对应关系。归一化将像素值从[0, 255]归一化到[-1, 1]与生成器输出层的Tanh激活函数匹配。4.2 训练循环的完整实现与技巧训练GAN需要精心设计训练循环。标准的流程是先更新判别器再更新生成器。import torch.optim as optim from torch.utils.tensorboard import SummaryWriter import numpy as np from skimage.metrics import peak_signal_noise_ratio as psnr_calc def train_epoch(generator, discriminator, train_loader, g_loss_fn, d_loss_fn, g_optimizer, d_optimizer, epoch, writer, device): generator.train() discriminator.train() total_g_loss 0 total_d_loss 0 for batch_idx, (lr_imgs, hr_imgs) in enumerate(train_loader): lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) # --------------------- # 训练判别器 (Discriminator) # --------------------- d_optimizer.zero_grad() # 生成假的高清图 with torch.no_grad(): # 生成器不参与判别器梯度计算 fake_hr generator(lr_imgs) # 判别器对真实图像和生成图像的判断 real_preds discriminator(hr_imgs) fake_preds discriminator(fake_imgs.detach()) # 使用.detach()切断与生成器的梯度联系 # 判别器损失希望将真实图判为1生成图判为0 real_labels torch.ones_like(real_preds, devicedevice) fake_labels torch.zeros_like(fake_preds, devicedevice) d_loss_real d_loss_fn(real_preds, real_labels) d_loss_fake d_loss_fn(fake_preds, fake_labels) d_loss (d_loss_real d_loss_fake) / 2 d_loss.backward() d_optimizer.step() # --------------------- # 训练生成器 (Generator) # --------------------- g_optimizer.zero_grad() # 再次生成图像这次需要梯度 fake_imgs generator(lr_imgs) fake_preds_for_g discriminator(fake_imgs) # 判别器对生成图的判断 # 生成器总损失 g_loss, mse_l, percep_l, adv_l g_loss_fn(fake_imgs, hr_imgs, fake_preds_for_g) g_loss.backward() g_optimizer.step() total_g_loss g_loss.item() total_d_loss d_loss.item() # 每N个batch记录一次日志 if batch_idx % 100 0: print(fEpoch: {epoch} [{batch_idx}/{len(train_loader)}] fD Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f} f(MSE: {mse_l.item():.4f}, Percep: {percep_l.item():.4f}, Adv: {adv_l.item():.4f})) # 记录到TensorBoard step epoch * len(train_loader) batch_idx writer.add_scalar(Loss/Discriminator, d_loss.item(), step) writer.add_scalar(Loss/Generator_Total, g_loss.item(), step) writer.add_scalar(Loss/Generator_MSE, mse_l.item(), step) writer.add_scalar(Loss/Generator_Perceptual, percep_l.item(), step) writer.add_scalar(Loss/Generator_Adversarial, adv_l.item(), step) avg_g_loss total_g_loss / len(train_loader) avg_d_loss total_d_loss / len(train_loader) return avg_g_loss, avg_d_loss训练技巧与参数设置优化器选择论文中使用Adam优化器。对于生成器和判别器通常使用相同的初始学习率如1e-4但beta参数可能不同。判别器有时会用稍小的学习率以防止它变得太强。g_optimizer optim.Adam(generator.parameters(), lr1e-4, betas(0.9, 0.999)) d_optimizer optim.Adam(discriminator.parameters(), lr1e-4, betas(0.9, 0.999))学习率调度使用MultiStepLR在训练到一定epoch时降低学习率有助于模型收敛到更好的局部最优。g_scheduler optim.lr_scheduler.MultiStepLR(g_optimizer, milestones[100, 200], gamma0.1) d_scheduler optim.lr_scheduler.MultiStepLR(d_optimizer, milestones[100, 200], gamma0.1)梯度裁剪对于判别器特别是训练初期梯度可能爆炸。加入梯度裁剪能增加稳定性。torch.nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm1.0)标签平滑一种正则化技巧在计算判别器损失时将真实标签从1稍微降低如0.9将生成标签从0稍微提高如0.1可以防止判别器过于自信有助于生成器学习。real_labels torch.full_like(real_preds, 0.9, devicedevice) # 标签平滑 fake_labels torch.full_like(fake_preds, 0.1, devicedevice)4.3 模型评估、保存与可视化评估指标PSNR/SSIM 在验证集上定期评估模型性能至关重要。PSNR峰值信噪比是超分辨率领域最常用的客观指标单位dB越高越好SSIM结构相似性则更符合人眼感知。def evaluate(generator, val_loader, device): generator.eval() total_psnr 0.0 total_ssim 0.0 with torch.no_grad(): for lr_imgs, hr_imgs in val_loader: lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) sr_imgs generator(lr_imgs) # 将图像从[-1,1]转换回[0,255]的uint8格式用于计算 sr_imgs_np (sr_imgs.cpu().numpy().transpose(0,2,3,1) * 127.5 127.5).astype(np.uint8) hr_imgs_np (hr_imgs.cpu().numpy().transpose(0,2,3,1) * 127.5 127.5).astype(np.uint8) for i in range(sr_imgs_np.shape[0]): psnr_val psnr_calc(hr_imgs_np[i], sr_imgs_np[i], data_range255) # ssim_val ssim(hr_imgs_np[i], sr_imgs_np[i], multichannelTrue, data_range255) total_psnr psnr_val # total_ssim ssim_val avg_psnr total_psnr / len(val_loader.dataset) # avg_ssim total_ssim / len(val_loader.dataset) return avg_psnr #, avg_ssim模型保存策略 我们实现两种保存方式定期保存Checkpoint保存模型权重、优化器状态、学习率调度器状态、当前epoch等。用于恢复训练。保存最优PSNR模型在验证集上监控PSNR只保存指标最高的模型。best_psnr 0.0 for epoch in range(start_epoch, num_epochs): train(...) val_psnr evaluate(...) # 保存checkpoint checkpoint { epoch: epoch, generator_state_dict: generator.state_dict(), discriminator_state_dict: discriminator.state_dict(), g_optimizer_state_dict: g_optimizer.state_dict(), d_optimizer_state_dict: d_optimizer.state_dict(), best_psnr: best_psnr, } torch.save(checkpoint, fcheckpoint_epoch_{epoch}.pth) # 保存最优模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(generator.state_dict(), fbest_generator_psnr_{val_psnr:.2f}.pth) print(fNew best PSNR: {val_psnr:.2f} dB, model saved.)训练曲线可视化 使用TensorBoard可以实时查看损失和PSNR曲线直观判断训练状态。writer SummaryWriter(log_dirruns/srgan_experiment_1) # 在训练循环中记录标量 writer.add_scalar(PSNR/val, val_psnr, epoch) # 还可以记录图像对比 if epoch % 10 0: writer.add_images(Val/HR, hr_imgs[:4], epoch) writer.add_images(Val/SR, sr_imgs[:4], epoch) writer.add_images(Val/LR, lr_imgs[:4], epoch)在命令行运行tensorboard --logdirruns即可在浏览器查看。5. 常见问题排查与调优经验训练SRGAN的过程很少一帆风顺以下是我在多次复现中总结的典型问题与解决方案。5.1 训练不稳定与模式崩溃现象判别器损失迅速降为0生成器损失飙升或者两者损失剧烈振荡。生成的图像全是噪声或重复的简单纹理。原因与对策判别器过强判别器学习太快导致生成器梯度消失梯度为0或很小。对策降低判别器的学习率减少判别器的更新频率例如每更新生成器2次才更新判别器1次对判别器使用梯度惩罚如WGAN-GP中的梯度惩罚项这是更现代、更稳定的方法。学习率过高这是新手最常见的问题。对策将初始学习率从1e-4降低到5e-5甚至1e-5试试。没有进行生成器预训练直接开始对抗训练对生成器要求太高。对策先用MSE损失仅内容损失单独训练生成器10-50个epoch得到一个基础超分模型再加载这个预训练权重加入判别器进行联合训练。使用WGAN-GP损失这是解决GAN训练不稳定的利器。它用Wasserstein距离代替JS散度并通过梯度惩罚Gradient Penalty来满足Lipschitz约束。在实践中将原始GAN的BCE损失换成WGAN-GP损失稳定性会有质的提升。5.2 生成图像模糊或伪影现象图像整体模糊缺乏纹理细节或者出现规则的棋盘格伪影Checkerboard Artifacts。原因与对策内容损失权重过高如果MSE损失权重占主导模型会倾向于输出所有可能结果的平均导致模糊。对策适当降低内容损失权重content_weight提高感知损失perceptual_weight和对抗损失adversarial_weight的权重引导模型生成更锐利的细节。棋盘格伪影这通常由上采样层如转置卷积引起。对策SRGAN已经使用了PixelShuffle基本避免了此问题。如果你修改了网络请确保上采样模块使用PixelShuffle或最近邻上采样卷积的组合。感知损失层选择不当原始SRGAN使用VGG19的conv5_4层之前的特征。如果使用太浅的层如conv3_3可能对纹理细节的约束不够使用太深的层可能过于关注高级语义而忽略局部纹理。对策可以尝试不同层的组合或使用更现代的感知损失如基于ResNet或LPIPSLearned Perceptual Image Patch Similarity的损失。5.3 PSNR指标不升反降现象随着训练进行视觉质量似乎有提升但验证集PSNR却下降了。原因与对策PSNR与感知质量的固有矛盾PSNR衡量的是像素级误差而GAN旨在优化感知质量。生成器为了产生更逼真的纹理可能会在像素位置上做一些“合理”的偏移这会降低PSNR。这是正常现象。你需要同时关注SSIM和进行人工主观评估。一个在PSNR上不是最高但视觉更清晰的模型往往是更成功的。过拟合模型在训练集上表现很好但在验证集上PSNR下降。对策使用更多的数据增强在生成器和判别器中加入Dropout层如果数据集不大可以尝试减小模型容量减少残差块数量。验证集处理不一致计算PSNR时图像的范围必须是[0, 255]。确保你的验证集图像在送入模型前和从模型输出后经过了与训练集完全相同的归一化和反归一化流程。一个常见的错误是归一化时用了错误的均值/标准差。5.4 显存不足OOM问题现象训练时出现CUDA out of memory错误。原因与对策Batch Size过大这是主因。对策减小batch_size。对于SRGANbatch_size16或32在11G显存的GPU上通常可行。如果必须用小batch可以累积梯度Gradient Accumulation每N个小batch进行一次参数更新相当于模拟大batch的效果。accumulation_steps 4 for i, (data, target) in enumerate(train_loader): ... loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()图像Patch过大减少训练时裁剪的HR patch大小如从96降到64。使用混合精度训练PyTorch的AMPAutomatic Mixed Precision自动使用半精度FP16进行计算可以显著减少显存占用并加速训练通常对最终精度影响很小。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): fake_imgs generator(lr_imgs) loss criterion(fake_imgs, hr_imgs) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5.5 项目代码使用与扩展建议拿到这份复现代码后你可以快速开始配置好环境下载DIV2K数据集运行python train.py --config configs/srgan_x4.yaml即可开始训练。配置文件里已经调好了大部分参数。使用预训练模型项目提供了在常用数据集上预训练的x2, x4, x8模型权重。你可以直接加载进行推理generator Generator(scale_factor4).to(device) generator.load_state_dict(torch.load(weights/best_generator_x4.pth, map_locationdevice)) generator.eval()在自己的数据上微调如果你想针对特定类型的图像如人脸、动漫、医学影像优化可以用预训练模型作为起点在自己的小数据集上用较低的学习率进行微调这会比从头训练快得多效果也更好。尝试改进模型这个项目是一个坚实的基础。你可以在此基础上轻松替换生成器为更先进的ESRGAN、Real-ESRGAN的架构或者将判别器损失换成WGAN-GP、Hinge Loss也可以尝试不同的感知损失网络。代码模块化的设计让这些实验变得很方便。训练一个高质量的SRGAN模型需要耐心和对细节的把控。从数据准备、损失权重的调整到训练策略的选择每一步都可能影响最终结果。这份详细的复现代码和指南希望能帮你绕过我当年踩过的那些坑更顺畅地进入图像超分辨率这个有趣而又充满挑战的领域。记住当看到第一张由你自己训练的模型生成的、细节丰富的超分辨率图片时那种成就感会让你觉得所有的调试都是值得的。本文还有配套的精品资源点击获取 PyTorch复现SRGAN：从理论到工程实践的超分辨率完整指南
简介本资源是图像超分辨率领域经典模型SRGAN的PyTorch完整复现工程面向计算机视觉方向的初学者与进阶研究者聚焦真实场景下的单图超分任务尤其适用于课程设计、科研复现实验及轻量级项目集成。压缩包共375个文件含10个核心Python脚本如train.py、test_image.py、draw_evaluation.py等、324张测试/验证用PNG图像涵盖barbara、lenna、baboon等标准基准图、3个最优PSNR模型权重x2/x4/x8倍率及训练过程产出的epochs、statistics、training_results等结构化结果目录整体大小231.09MB。已有694人学习下载资源配套CSDN技术博文详解实现逻辑与调参要点。用户可直接运行demo.py进行GT/Bicubic/SRGAN三栏可视化对比调用test_video.py处理视频序列并通过draw_evaluation.py一键生成Loss、PSNR、SSIM随Epoch变化曲线所有模块注释详尽、路径规范、即插即用。1. 项目概述从SRGAN到可复现的PyTorch工程如果你在图像处理领域摸爬滚打过一阵子肯定对“超分辨率”这个词不陌生。简单说就是把一张模糊的、像素低的小图通过算法“脑补”出细节变成一张清晰的高清大图。这技术从早期的插值算法到后来的深度学习模型发展速度飞快。而SRGAN可以说是这个领域里一个里程碑式的存在。它第一次把生成对抗网络GAN的思路引入超分辨率不再仅仅追求像素值上的接近比如PSNR指标而是开始关注人眼视觉感受生成纹理更真实、细节更丰富的图像。我最初接触SRGAN时是在读它的论文。理论很优美但真到了自己动手复现的时候才发现从论文到代码之间隔着一片“坑海”。官方没有提供完整的训练代码网上的开源实现质量参差不齐有的注释不清有的依赖老旧有的甚至关键步骤都有错误。训练过程更是玄学GAN本身就不稳定再加上超分辨率任务对细节的极致追求调参调得人头皮发麻。更别提还要自己写训练日志、画损失曲线、保存最优模型这些工程化的事情了一套流程下来精力全耗在搭轮子上反而没多少时间去理解模型本身了。所以当我决定重新梳理并实现一个SRGAN的PyTorch版本时目标就非常明确不仅要代码能跑通更要让它成为一个“教学级”和“生产级”兼备的项目。这意味着每一行核心代码都要有详尽的注释解释清楚“为什么这么做”训练过程要可视化让你能实时看到模型是学好了还是在“摆烂”还要自动保存PSNR指标最好的模型权重方便你直接拿来用或者继续微调。最终这个项目覆盖了x2、x4、x8三种常见的超分倍数你拿到手就能在自己的数据上跑起来或者以此为蓝本去实现更复杂的变体。2. 核心思路与方案选型为什么是PyTorch SRGAN2.1 模型架构的再审视生成器与判别器的博弈SRGAN的核心思想是“对抗训练”。它包含两个网络生成器Generator它的任务是把低分辨率LR图像“想象”成高分辨率HR图像。在SRGAN中生成器的主体是一个深度残差网络。为什么用残差结构因为超分任务可以看作是在LR图像基础上添加高频细节残差残差连接能有效缓解深层网络的梯度消失问题让网络专注于学习LR到HR的“差异部分”训练起来更稳定、更快。判别器Discriminator它是一个二分类器任务就是判断一张图像是“真实的”高清图来自数据集还是“伪造的”由生成器产生的高清图。它的存在就像一位严厉的考官不断逼迫生成器去生成以假乱真的图像。两者的关系是动态博弈生成器努力骗过判别器判别器努力识破生成器。这个博弈过程最终驱使生成器产生的图像不仅在像素值上接近真实在纹理、边缘等视觉感知上也更加逼真。这与传统只使用MSE均方误差损失的方法有本质区别。MSE损失容易导致结果过于平滑丢失纹理因为它是“平均主义者”而GAN的对抗损失鼓励生成“局部逼真”的细节。注意原始SRGAN论文中生成器的损失函数是内容损失Content Loss、**对抗损失Adversarial Loss和感知损失Perceptual Loss基于VGG特征**的加权和。其中感知损失是让生成图像和真实图像在深层特征空间而非像素空间上接近这是提升视觉质量的关键。我们的复现严格遵循了这一设计。2.2 为什么选择PyTorch作为实现框架面对TensorFlow、PyTorch乃至JAX等框架我选择了PyTorch这是基于多年一线开发经验的考量动态图优先PyTorch的 eager execution 模式让调试变得异常直观。你可以在任意地方插入print或使用调试器查看张量的形状和值这对于理解GAN这种复杂、动态的训练过程至关重要。想象一下当判别器损失突然变成NaN时你能快速定位到是梯度爆炸还是数据出了问题。生态与社区PyTorch在学术研究和工业界的普及度已毋庸置疑。这意味着你遇到的几乎所有问题都能在社区找到讨论和解决方案。从最新的模型实现如ESRGAN、Real-ESRGAN到各种数据加载、可视化工具PyTorch的生态是最丰富的。简洁的API设计torch.nn.Module和torch.optim的设计非常Pythonic定义模型和训练循环的逻辑清晰直白。自定义层、损失函数都轻而易举这让实现SRGAN中特定的残差块Residual Block和感知损失变得很顺畅。部署友好虽然本项目侧重于训练和复现但PyTorch模型通过TorchScript或ONNX可以相对方便地转换为后续移动端或服务端部署留下了可能性。相比之下虽然TensorFlow 2.x也改善了易用性但其历史包袱和某些API设计仍显繁琐。对于这样一个以清晰、教育为目的的复现项目PyTorch是更自然的选择。2.3 项目特色与解决的问题市面上有很多SRGAN代码那这个复现项目有什么不同它主要解决了以下几个痛点注释的深度注释不仅仅是解释“这行代码在做什么”What更重要的是解释“为什么要这么做”Why。例如在初始化权重时会说明为何使用kaiming_normal_初始化而不是xavier_uniform_在训练循环中会解释为何要先更新判别器再更新生成器。训练过程透明化训练GAN就像驾驶一架仪表盘不全的飞机。本项目集成了TensorBoard或Matplotlib实时绘图将生成器损失、判别器损失、PSNR、SSIM等关键指标的变化曲线实时可视化。你能一眼看出模型是否收敛是否发生了模式崩溃Mode Collapse。模型管理的自动化我们不止在每一个epoch结束时保存模型。代码会自动监控验证集上的PSNR指标只保存达到历史最优PSNR的那个模型权重。同时也会保存最新的模型作为checkpoint防止训练意外中断。这意味着你训练结束后直接就能拿到“表现最好”的模型而不是“最后一代”可能已经过拟合的模型。开箱即用的配置提供了x2, x4, x8三种上采样倍数的预训练模型或训练脚本配置。数据集处理、数据增强翻转、旋转、学习率调度MultiStepLR等常用技巧都已集成并配有详细的配置文件如YAML你只需要修改数据路径和几个参数就能跑起来。3. 环境搭建与核心代码解析3.1 从零开始PyTorch与CUDA环境配置工欲善其事必先利其器。一个稳定、版本匹配的深度学习环境是成功的第一步。很多人在这里踩坑问题多半出在CUDA、cuDNN和PyTorch的版本不匹配上。步骤一确认你的GPU和CUDA驱动打开命令行输入nvidia-smi。记下最上方显示的CUDA Version例如12.4。这是你的驱动支持的最高CUDA运行时版本但PyTorch需要安装与之兼容的CUDA Toolkit。步骤二使用Conda创建独立环境强烈建议使用Conda管理环境避免包冲突。conda create -n srgan_pytorch python3.9 conda activate srgan_pytorch选择Python 3.9是因为它在稳定性和对新旧包的兼容性上比较平衡。步骤三安装匹配的PyTorch这是最关键的一步。不要直接用pip install torch。前往 PyTorch官网 根据你的系统、Conda环境以及上一步查到的CUDA驱动版本选择对应的安装命令。 例如对于CUDA 12.1你可能会看到pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121对于没有NVIDIA GPU的机器则安装CPU版本。步骤四安装项目依赖在项目根目录下通常有一个requirements.txt文件。pip install -r requirements.txt典型的依赖包括opencv-python图像处理、pillow图像读取、tensorboard或tensorboardX可视化、matplotlib绘图、scikit-image计算PSNR/SSIM、tqdm进度条等。实操心得我习惯在安装完PyTorch后立刻写一个测试脚本验证GPU是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()返回False大概率是CUDA版本不匹配需要重新安装PyTorch或升级NVIDIA驱动。3.2 生成器网络结构深度拆解SRGAN的生成器是一个深层的残差网络。我们来逐层解析其PyTorch实现的关键部分。import torch.nn as nn import torch.nn.functional as F class ResidualBlock(nn.Module): 残差块两个卷积层 跳跃连接用于学习高频细节残差 def __init__(self, channels): super(ResidualBlock, self).__init__() self.conv1 nn.Conv2d(channels, channels, kernel_size3, padding1, biasFalse) self.bn1 nn.BatchNorm2d(channels) self.prelu nn.PReLU() # 使用PReLU而非ReLU允许小的负斜率可能提升性能 self.conv2 nn.Conv2d(channels, channels, kernel_size3, padding1, biasFalse) self.bn2 nn.BatchNorm2d(channels) def forward(self, x): residual x out self.conv1(x) out self.bn1(out) out self.prelu(out) out self.conv2(out) out self.bn2(out) # 残差连接将输入x加到输出上 out out residual return out class UpsampleBlock(nn.Module): 上采样块用于将特征图空间分辨率扩大2倍 def __init__(self, in_channels, up_scale): super(UpsampleBlock, self).__init__() # 使用PixelShuffle进行高效上采样 self.conv nn.Conv2d(in_channels, in_channels * (up_scale ** 2), kernel_size3, padding1) self.pixel_shuffle nn.PixelShuffle(up_scale) # 重排像素实现上采样 self.prelu nn.PReLU() def forward(self, x): out self.conv(x) out self.pixel_shuffle(out) out self.prelu(out) return out class Generator(nn.Module): SRGAN生成器网络 def __init__(self, scale_factor4, num_residual_blocks16): super(Generator, self).__init__() # 初始卷积层将输入图像映射到特征空间 self.conv1 nn.Sequential( nn.Conv2d(3, 64, kernel_size9, padding4), nn.PReLU() ) # 一系列残差块用于深度特征提取和学习残差 residual_blocks [] for _ in range(num_residual_blocks): residual_blocks.append(ResidualBlock(64)) self.residual_blocks nn.Sequential(*residual_blocks) # 残差块后的卷积层 self.conv2 nn.Sequential( nn.Conv2d(64, 64, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(64) ) # 上采样模块根据放大倍数决定上采样次数 upsample_blocks [] # 计算需要几次2倍上采样。例如 scale_factor4 - num_up2 num_up int(math.log2(scale_factor)) for _ in range(num_up): upsample_blocks.append(UpsampleBlock(64, up_scale2)) self.upsample_blocks nn.Sequential(*upsample_blocks) # 最终输出层将特征映射回RGB图像 self.conv3 nn.Conv2d(64, 3, kernel_size9, padding4) def forward(self, x): # 初始特征提取 out1 self.conv1(x) # 通过残差块 out self.residual_blocks(out1) # 残差连接大残差将初始特征与残差块输出相加 out self.conv2(out) out out out1 # 上采样 out self.upsample_blocks(out) # 最终输出使用Tanh将值约束到[-1, 1]与归一化后的输入数据匹配 out torch.tanh(self.conv3(out)) return out关键点解析PixelShuffle上采样这是SRGAN论文中的关键。传统上采样如转置卷积Deconv容易产生棋盘伪影Checkerboard Artifacts。PixelShuffle通过卷积增加通道数然后进行像素重排来扩大空间尺寸能生成更平滑的结果。例如将(C, H, W)的特征图通过卷积变为(C*4, H, W)再通过PixelShuffle(2)重排为(C, H*2, W*2)。残差连接的作用代码中有两处残差连接。一是在每个ResidualBlock内部保证梯度流动二是在所有残差块之后out out out1这是一个“大残差”连接让网络更容易学习到LR与HR之间的残差即细节信息而非学习完整的图像映射这大大降低了学习难度。输出激活函数Tanh输入图像在预处理时通常被归一化到[-1, 1]范围因此输出层使用Tanh将像素值约束到同一范围便于计算损失。3.3 判别器与损失函数的设计判别器是一个相对标准的CNN分类器但针对图像任务做了一些设计。class Discriminator(nn.Module): SRGAN判别器网络一个二分类CNN def __init__(self): super(Discriminator, self).__init__() self.net nn.Sequential( # 输入: [3, 96, 96] (假设HR patch大小) nn.Conv2d(3, 64, kernel_size3, padding1), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(64, 64, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(64), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸减半 nn.Conv2d(64, 128, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(128, 128, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸再减半 nn.Conv2d(128, 256, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(256, 256, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplaceTrue), # 特征图尺寸再减半 nn.Conv2d(256, 512, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(512, 512, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplaceTrue), # 此时特征图已很小 nn.AdaptiveAvgPool2d(1), # 全局平均池化得到一个512维的向量 nn.Conv2d(512, 1024, kernel_size1), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv2d(1024, 1, kernel_size1), # 输出一个标量通过Sigmoid表示“真”的概率 ) def forward(self, x): out self.net(x) return out.view(-1, 1) # 展平为 [batch_size, 1]损失函数的组合是SRGAN的灵魂。生成器的总损失是三项的加权和import torch import torch.nn as nn from torchvision.models import vgg19 class GeneratorLoss(nn.Module): def __init__(self, content_weight1e-2, adversarial_weight1e-3): super(GeneratorLoss, self).__init__() # 1. 内容损失 (MSE Loss) - 保证像素级相似性 self.mse_loss nn.MSELoss() # 2. 对抗损失 (BCE Loss) - 鼓励生成器骗过判别器 self.adversarial_loss nn.BCEWithLogitsLoss() # 注意判别器输出未经过Sigmoid故用WithLogits # 3. 感知损失 (Perceptual Loss) - 基于VGG特征保证高级语义相似性 vgg vgg19(pretrainedTrue).features[:36] # 取VGG19的前36层到conv5_4 for param in vgg.parameters(): param.requires_grad False # 冻结VGG参数 self.vgg vgg self.vgg_loss nn.MSELoss() self.content_weight content_weight self.adversarial_weight adversarial_weight def forward(self, sr_images, hr_images, discriminator_fake_output): # 计算MSE Loss mse_l self.mse_loss(sr_images, hr_images) # 计算对抗Loss目标是让判别器认为生成图像为“真”(label1) adversarial_l self.adversarial_loss(discriminator_fake_output, torch.ones_like(discriminator_fake_output)) # 计算感知Loss提取VGG特征 sr_features self.vgg(sr_images) hr_features self.vgg(hr_images) perceptual_l self.vgg_loss(sr_features, hr_features) # 加权总损失 total_loss mse_l self.content_weight * perceptual_l self.adversarial_weight * adversarial_l return total_loss, mse_l, perceptual_l, adversarial_l权重选择经验content_weight和adversarial_weight的平衡至关重要。对抗权重太大训练容易不稳定图像可能出现奇怪纹理太小则退化回普通MSE模型图像平滑。论文中分别使用1e-2和1e-3这是一个很好的起点。在实际训练中我通常会先以较大的内容损失权重如1.0和极小的对抗损失权重如1e-4预训练生成器几十个epoch让生成器先学会一个“基础版”的超分然后再调高对抗权重进行“对抗精炼”这样训练更稳定。4. 数据准备与训练流程实战4.1 数据集处理与高效数据加载高质量的数据是训练成功的一半。对于超分辨率任务你需要成对的低分辨率LR和高分辨率HR图像。通常有两种获取方式使用现成数据集如DIV2K、Flickr2K等它们已经提供了HR图像。你需要在线下用双三次下采样Bicubic Downsampling生成对应的LR图像。关键点下采样的插值算法必须与你的评估标准如PSNR所使用的算法一致通常是MATLAB的imresize函数。为了复现论文结果我强烈建议使用论文作者提供的脚本或opencv设置INTER_CUBIC来生成LR图并确保scale_factor准确。自定义数据如果你有自己的高清图库处理流程同上。数据加载器的构建from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as transforms import os class SRDataset(Dataset): def __init__(self, hr_dir, lr_dir, patch_size96, scale4, is_trainTrue): self.hr_dir hr_dir self.lr_dir lr_dir self.scale scale self.patch_size patch_size self.is_train is_train self.hr_images sorted([os.path.join(hr_dir, x) for x in os.listdir(hr_dir) if x.endswith((.png, .jpg))]) self.lr_images sorted([os.path.join(lr_dir, x) for x in os.listdir(lr_dir) if x.endswith((.png, .jpg))]) # 训练和验证/测试的数据转换不同 if is_train: self.transform transforms.Compose([ transforms.RandomCrop(patch_size), # 随机裁剪patch transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) # 归一化到[-1,1] ]) else: self.transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) ]) def __len__(self): return len(self.hr_images) def __getitem__(self, idx): hr_img Image.open(self.hr_images[idx]).convert(RGB) lr_img Image.open(self.lr_images[idx]).convert(RGB) # 确保LR和HR尺寸匹配 (HR LR * scale) lr_width, lr_height lr_img.size hr_width, hr_height hr_img.size assert hr_width lr_width * self.scale and hr_height lr_height * self.scale, \ f尺寸不匹配: LR {lr_img.size}, HR {hr_img.size}, scale {self.scale} 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_cropped transforms.functional.crop(hr_img, i, j, h, w) # 对LR进行对应区域的裁剪 (i//scale, j//scale, h//scale, w//scale) lr_img_cropped transforms.functional.crop(lr_img, i//self.scale, j//self.scale, h//self.scale, w//self.scale) hr_img hr_img_cropped lr_img lr_img_cropped hr_tensor self.transform(hr_img) lr_tensor self.transform(lr_img) return lr_tensor, hr_tensor重要细节Patch训练由于高清图像很大无法整张送入网络。通常随机裁剪出固定大小如96x96的HR patch并对应裁剪出LR patch进行训练。这增加了数据的多样性。数据增强随机水平/垂直翻转是基本操作有时还会加入随机旋转90, 180, 270度来进一步增强。注意增强操作应同时应用于配对的LR和HR patch保持空间对应关系。归一化将像素值从[0, 255]归一化到[-1, 1]与生成器输出层的Tanh激活函数匹配。4.2 训练循环的完整实现与技巧训练GAN需要精心设计训练循环。标准的流程是先更新判别器再更新生成器。import torch.optim as optim from torch.utils.tensorboard import SummaryWriter import numpy as np from skimage.metrics import peak_signal_noise_ratio as psnr_calc def train_epoch(generator, discriminator, train_loader, g_loss_fn, d_loss_fn, g_optimizer, d_optimizer, epoch, writer, device): generator.train() discriminator.train() total_g_loss 0 total_d_loss 0 for batch_idx, (lr_imgs, hr_imgs) in enumerate(train_loader): lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) # --------------------- # 训练判别器 (Discriminator) # --------------------- d_optimizer.zero_grad() # 生成假的高清图 with torch.no_grad(): # 生成器不参与判别器梯度计算 fake_hr generator(lr_imgs) # 判别器对真实图像和生成图像的判断 real_preds discriminator(hr_imgs) fake_preds discriminator(fake_imgs.detach()) # 使用.detach()切断与生成器的梯度联系 # 判别器损失希望将真实图判为1生成图判为0 real_labels torch.ones_like(real_preds, devicedevice) fake_labels torch.zeros_like(fake_preds, devicedevice) d_loss_real d_loss_fn(real_preds, real_labels) d_loss_fake d_loss_fn(fake_preds, fake_labels) d_loss (d_loss_real d_loss_fake) / 2 d_loss.backward() d_optimizer.step() # --------------------- # 训练生成器 (Generator) # --------------------- g_optimizer.zero_grad() # 再次生成图像这次需要梯度 fake_imgs generator(lr_imgs) fake_preds_for_g discriminator(fake_imgs) # 判别器对生成图的判断 # 生成器总损失 g_loss, mse_l, percep_l, adv_l g_loss_fn(fake_imgs, hr_imgs, fake_preds_for_g) g_loss.backward() g_optimizer.step() total_g_loss g_loss.item() total_d_loss d_loss.item() # 每N个batch记录一次日志 if batch_idx % 100 0: print(fEpoch: {epoch} [{batch_idx}/{len(train_loader)}] fD Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f} f(MSE: {mse_l.item():.4f}, Percep: {percep_l.item():.4f}, Adv: {adv_l.item():.4f})) # 记录到TensorBoard step epoch * len(train_loader) batch_idx writer.add_scalar(Loss/Discriminator, d_loss.item(), step) writer.add_scalar(Loss/Generator_Total, g_loss.item(), step) writer.add_scalar(Loss/Generator_MSE, mse_l.item(), step) writer.add_scalar(Loss/Generator_Perceptual, percep_l.item(), step) writer.add_scalar(Loss/Generator_Adversarial, adv_l.item(), step) avg_g_loss total_g_loss / len(train_loader) avg_d_loss total_d_loss / len(train_loader) return avg_g_loss, avg_d_loss训练技巧与参数设置优化器选择论文中使用Adam优化器。对于生成器和判别器通常使用相同的初始学习率如1e-4但beta参数可能不同。判别器有时会用稍小的学习率以防止它变得太强。g_optimizer optim.Adam(generator.parameters(), lr1e-4, betas(0.9, 0.999)) d_optimizer optim.Adam(discriminator.parameters(), lr1e-4, betas(0.9, 0.999))学习率调度使用MultiStepLR在训练到一定epoch时降低学习率有助于模型收敛到更好的局部最优。g_scheduler optim.lr_scheduler.MultiStepLR(g_optimizer, milestones[100, 200], gamma0.1) d_scheduler optim.lr_scheduler.MultiStepLR(d_optimizer, milestones[100, 200], gamma0.1)梯度裁剪对于判别器特别是训练初期梯度可能爆炸。加入梯度裁剪能增加稳定性。torch.nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm1.0)标签平滑一种正则化技巧在计算判别器损失时将真实标签从1稍微降低如0.9将生成标签从0稍微提高如0.1可以防止判别器过于自信有助于生成器学习。real_labels torch.full_like(real_preds, 0.9, devicedevice) # 标签平滑 fake_labels torch.full_like(fake_preds, 0.1, devicedevice)4.3 模型评估、保存与可视化评估指标PSNR/SSIM 在验证集上定期评估模型性能至关重要。PSNR峰值信噪比是超分辨率领域最常用的客观指标单位dB越高越好SSIM结构相似性则更符合人眼感知。def evaluate(generator, val_loader, device): generator.eval() total_psnr 0.0 total_ssim 0.0 with torch.no_grad(): for lr_imgs, hr_imgs in val_loader: lr_imgs, hr_imgs lr_imgs.to(device), hr_imgs.to(device) sr_imgs generator(lr_imgs) # 将图像从[-1,1]转换回[0,255]的uint8格式用于计算 sr_imgs_np (sr_imgs.cpu().numpy().transpose(0,2,3,1) * 127.5 127.5).astype(np.uint8) hr_imgs_np (hr_imgs.cpu().numpy().transpose(0,2,3,1) * 127.5 127.5).astype(np.uint8) for i in range(sr_imgs_np.shape[0]): psnr_val psnr_calc(hr_imgs_np[i], sr_imgs_np[i], data_range255) # ssim_val ssim(hr_imgs_np[i], sr_imgs_np[i], multichannelTrue, data_range255) total_psnr psnr_val # total_ssim ssim_val avg_psnr total_psnr / len(val_loader.dataset) # avg_ssim total_ssim / len(val_loader.dataset) return avg_psnr #, avg_ssim模型保存策略 我们实现两种保存方式定期保存Checkpoint保存模型权重、优化器状态、学习率调度器状态、当前epoch等。用于恢复训练。保存最优PSNR模型在验证集上监控PSNR只保存指标最高的模型。best_psnr 0.0 for epoch in range(start_epoch, num_epochs): train(...) val_psnr evaluate(...) # 保存checkpoint checkpoint { epoch: epoch, generator_state_dict: generator.state_dict(), discriminator_state_dict: discriminator.state_dict(), g_optimizer_state_dict: g_optimizer.state_dict(), d_optimizer_state_dict: d_optimizer.state_dict(), best_psnr: best_psnr, } torch.save(checkpoint, fcheckpoint_epoch_{epoch}.pth) # 保存最优模型 if val_psnr best_psnr: best_psnr val_psnr torch.save(generator.state_dict(), fbest_generator_psnr_{val_psnr:.2f}.pth) print(fNew best PSNR: {val_psnr:.2f} dB, model saved.)训练曲线可视化 使用TensorBoard可以实时查看损失和PSNR曲线直观判断训练状态。writer SummaryWriter(log_dirruns/srgan_experiment_1) # 在训练循环中记录标量 writer.add_scalar(PSNR/val, val_psnr, epoch) # 还可以记录图像对比 if epoch % 10 0: writer.add_images(Val/HR, hr_imgs[:4], epoch) writer.add_images(Val/SR, sr_imgs[:4], epoch) writer.add_images(Val/LR, lr_imgs[:4], epoch)在命令行运行tensorboard --logdirruns即可在浏览器查看。5. 常见问题排查与调优经验训练SRGAN的过程很少一帆风顺以下是我在多次复现中总结的典型问题与解决方案。5.1 训练不稳定与模式崩溃现象判别器损失迅速降为0生成器损失飙升或者两者损失剧烈振荡。生成的图像全是噪声或重复的简单纹理。原因与对策判别器过强判别器学习太快导致生成器梯度消失梯度为0或很小。对策降低判别器的学习率减少判别器的更新频率例如每更新生成器2次才更新判别器1次对判别器使用梯度惩罚如WGAN-GP中的梯度惩罚项这是更现代、更稳定的方法。学习率过高这是新手最常见的问题。对策将初始学习率从1e-4降低到5e-5甚至1e-5试试。没有进行生成器预训练直接开始对抗训练对生成器要求太高。对策先用MSE损失仅内容损失单独训练生成器10-50个epoch得到一个基础超分模型再加载这个预训练权重加入判别器进行联合训练。使用WGAN-GP损失这是解决GAN训练不稳定的利器。它用Wasserstein距离代替JS散度并通过梯度惩罚Gradient Penalty来满足Lipschitz约束。在实践中将原始GAN的BCE损失换成WGAN-GP损失稳定性会有质的提升。5.2 生成图像模糊或伪影现象图像整体模糊缺乏纹理细节或者出现规则的棋盘格伪影Checkerboard Artifacts。原因与对策内容损失权重过高如果MSE损失权重占主导模型会倾向于输出所有可能结果的平均导致模糊。对策适当降低内容损失权重content_weight提高感知损失perceptual_weight和对抗损失adversarial_weight的权重引导模型生成更锐利的细节。棋盘格伪影这通常由上采样层如转置卷积引起。对策SRGAN已经使用了PixelShuffle基本避免了此问题。如果你修改了网络请确保上采样模块使用PixelShuffle或最近邻上采样卷积的组合。感知损失层选择不当原始SRGAN使用VGG19的conv5_4层之前的特征。如果使用太浅的层如conv3_3可能对纹理细节的约束不够使用太深的层可能过于关注高级语义而忽略局部纹理。对策可以尝试不同层的组合或使用更现代的感知损失如基于ResNet或LPIPSLearned Perceptual Image Patch Similarity的损失。5.3 PSNR指标不升反降现象随着训练进行视觉质量似乎有提升但验证集PSNR却下降了。原因与对策PSNR与感知质量的固有矛盾PSNR衡量的是像素级误差而GAN旨在优化感知质量。生成器为了产生更逼真的纹理可能会在像素位置上做一些“合理”的偏移这会降低PSNR。这是正常现象。你需要同时关注SSIM和进行人工主观评估。一个在PSNR上不是最高但视觉更清晰的模型往往是更成功的。过拟合模型在训练集上表现很好但在验证集上PSNR下降。对策使用更多的数据增强在生成器和判别器中加入Dropout层如果数据集不大可以尝试减小模型容量减少残差块数量。验证集处理不一致计算PSNR时图像的范围必须是[0, 255]。确保你的验证集图像在送入模型前和从模型输出后经过了与训练集完全相同的归一化和反归一化流程。一个常见的错误是归一化时用了错误的均值/标准差。5.4 显存不足OOM问题现象训练时出现CUDA out of memory错误。原因与对策Batch Size过大这是主因。对策减小batch_size。对于SRGANbatch_size16或32在11G显存的GPU上通常可行。如果必须用小batch可以累积梯度Gradient Accumulation每N个小batch进行一次参数更新相当于模拟大batch的效果。accumulation_steps 4 for i, (data, target) in enumerate(train_loader): ... loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()图像Patch过大减少训练时裁剪的HR patch大小如从96降到64。使用混合精度训练PyTorch的AMPAutomatic Mixed Precision自动使用半精度FP16进行计算可以显著减少显存占用并加速训练通常对最终精度影响很小。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): fake_imgs generator(lr_imgs) loss criterion(fake_imgs, hr_imgs) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5.5 项目代码使用与扩展建议拿到这份复现代码后你可以快速开始配置好环境下载DIV2K数据集运行python train.py --config configs/srgan_x4.yaml即可开始训练。配置文件里已经调好了大部分参数。使用预训练模型项目提供了在常用数据集上预训练的x2, x4, x8模型权重。你可以直接加载进行推理generator Generator(scale_factor4).to(device) generator.load_state_dict(torch.load(weights/best_generator_x4.pth, map_locationdevice)) generator.eval()在自己的数据上微调如果你想针对特定类型的图像如人脸、动漫、医学影像优化可以用预训练模型作为起点在自己的小数据集上用较低的学习率进行微调这会比从头训练快得多效果也更好。尝试改进模型这个项目是一个坚实的基础。你可以在此基础上轻松替换生成器为更先进的ESRGAN、Real-ESRGAN的架构或者将判别器损失换成WGAN-GP、Hinge Loss也可以尝试不同的感知损失网络。代码模块化的设计让这些实验变得很方便。训练一个高质量的SRGAN模型需要耐心和对细节的把控。从数据准备、损失权重的调整到训练策略的选择每一步都可能影响最终结果。这份详细的复现代码和指南希望能帮你绕过我当年踩过的那些坑更顺畅地进入图像超分辨率这个有趣而又充满挑战的领域。记住当看到第一张由你自己训练的模型生成的、细节丰富的超分辨率图片时那种成就感会让你觉得所有的调试都是值得的。本文还有配套的精品资源点击获取