076、SwinIR实战：基于Swin Transformer的图像超分辨率重建——从原理到代码逐行解析
去年年底有个项目，客户要求把监控视频里人脸从64x64放大到256x256，还要保留眉毛纹理。我一开始扔了个EDSR上去，结果放大后脸像打了马赛克，客户直接甩了句“这跟插值有啥区别”。后来换成SwinIR，虽然训练慢了点，但输出结果连眼睫毛都清晰可见。今天就把这个模型的实现细节掰开揉碎讲清楚，重点放在那些容易翻车的代码细节上。
为什么SwinIR能吊打传统CNN超分模型
先说说SwinIR的核心思路。传统SRCNN、EDSR这些模型用卷积核堆叠感受野，但卷积核天然存在两个问题：一是局部性太强，远处像素关联需要层层传递；二是下采样会丢失细节。SwinIR把Transformer那套自注意力机制搬过来，但直接套用ViT又不行——图像分辨率高时计算量爆炸。Swin Transformer的解决方案是：把图像切成不重叠的窗口，在每个窗口内算自注意力，再通过移位窗口实现跨窗口信息交互。
具体到超分任务，SwinIR做了三件事：浅层特征提取、深层特征提取（核心是多个Swin Transformer Block堆叠）、上采样重建。注意这里有个坑——浅层特征提取用的是卷积而不是Patch Embedding，因为卷积对高频细节更敏感，直接Patch Embedding会丢失边缘信息。
代码实现：从数据加载到模型推理
先看模型整体结构。我习惯把SwinIR拆成三个模块，调试时逐块验证输出shape。
classSwinIR(nn.Module):def__init__(self,upscale=4,img_size=64,window_size=8,embed_dim=180,depths=[6,6,6,6],num_heads=[6,6,6,6]):super().__init__()# 浅层特征提取：一个卷积层，别用BatchNorm，会破坏图像连续性self.conv_first=nn.Conv2d(3,embed_dim,3,1,1)# 深层特征提取：堆叠Swin Transformer Blockself.layers=nn.ModuleList()foriinrange(len(depths)):layer=SwinTransformerLayer(dim=embed_dim,depth=depths[i],num_heads=num_heads[i],window_size=window_size,img_size=img_size//(2**i)ifi>0elseimg_size# 这里踩过坑：下采样后窗口大小要对应调整)self.layers.append(layer)# 上采样部分：PixelShuffle + 卷积，别用反卷积，棋盘效应太明显self.upsample=nn.Sequential(nn.Conv2d(embed_dim,embed_dim*upscale*upscale,3,1,1),nn.PixelShuffle(upscale),nn.Conv2d(embed_dim,3,3,1,1))这里有个关键点：img_size参数。很多人直接写死64，但实际输入可能是任意尺寸。我后来改成动态计算：self.register_buffer('window_size', torch.tensor(window_size))，推理时根据输入尺寸自动调整窗口划分。
Swin Transformer Block：窗口注意力机制的实现细节
这是整个模型最绕的地方。每个Block包含两个子模块：W-MSA（窗口多头自注意力）和SW-MSA（移位窗口多头自注意力）。W-MSA把特征图切成不重叠的窗口，每个窗口内算自注意力；SW-MSA则把窗口向右下角偏移半个窗口大小，实现跨窗口交互。
classWindowAttention(nn.Module):def__init__(self,dim,num_heads,window_size):super().__init__()self.num_heads=num_heads self.window_size=window_size self.scale=(dim//num_heads)**-0.5# 相对位置偏置表：这是Swin的核心创新点# 别写成nn.Parameter(torch.zeros(2*window_size-1, 2*window_size-1))# 正确写法：每个head独立的位置偏置self.relative_position_bias_table=nn.Parameter(torch.zeros((2*window_size[0]-1)*(2*window_size[1]-1),num_heads))# 计算相对位置索引coords_h=torch.arange(window_size[0])coords_w=torch.arange(window_size[1])coords=torch.stack(torch.meshgrid([coords_h,coords_w]))# 2, Wh, Wwcoords_flatten=torch.flatten(coords,1)# 2, Wh*Wwrelative_coords=coords_flatten[:,:,None]-coords_flatten[:,None,:]# 2, Wh*Ww, Wh*Wwrelative_coords=relative_coords.permute(1,2,0).contiguous()# Wh*Ww, Wh*Ww, 2relative_coords[:,:,0]+=window_size[0]-1relative_coords[:,:,1]+=window_size[1]-1relative_coords[:,:,0]*=2*window_size[1]-1relative_position_index=relative_coords.sum(-1)# Wh*Ww, Wh*Wwself.register_buffer('relative_position_index',relative_position_index)注意这个相对位置偏置的计算。很多人直接复制官方代码但没理解为什么这么算。简单说：窗口内每个像素对(pixel_i, pixel_j)的相对位置坐标范围是[-(W-1), W-1]，需要映射到[0, 2W-2]作为索引。乘以(2W-1)是为了把二维坐标压成一维索引。这里如果窗口大小是8，索引表大小就是15x15=225，每个位置对应一个num_heads维的向量。
移位窗口的实现：torch.roll的坑
SW-MSA需要把特征图循环移位，让窗口边界处的像素能互相看到。官方实现用torch.roll，但这里有个隐藏问题：
defshift_window(x,shift_size,window_size):# 别这样写：torch.roll(x, shifts=(-shift_size, -shift_size), dims=(2,3))# 这样会导致边界像素循环到另一边，破坏空间连续性# 正确做法：先pad再roll，或者用mask机制ifshift_size>0:x=torch.roll(x,shifts=(-shift_size,-shift_size),dims=(2,3))# 生成attention mask，屏蔽掉不该交互的像素mask=generate_mask(x.shape[2],x.shape[3],window_size,shift_size)returnx,maskelse:returnx,Nonetorch.roll会让左上角的像素跑到右下角，如果不加mask，这些本不相邻的像素会错误地计算注意力。正确的做法是生成一个mask矩阵，在计算softmax前把非法位置的注意力值设为负无穷。这个mask的生成逻辑比较绕，我贴个简化版：
defgenerate_mask(H,W,window_size,shift_size):img_mask=torch.zeros((1,H,W,1))h_slices=(slice(0,-window_size),slice(-window_size,-shift_size),slice(-shift_size,None))w_slices=(slice(0,-window_size),slice(-window_size,-shift_size),slice(-shift_size,None))cnt=0forhinh_slices:forwinw_slices:img_mask[:,h,w,:]=cnt cnt+=1# 然后按窗口划分，每个窗口内的mask值不同，相同值的像素才能交互# 具体实现参考官方代码的window_partition函数上采样模块：为什么PixelShuffle比反卷积好
很多人喜欢用nn.ConvTranspose2d做上采样，但反卷积会产生棋盘格伪影。PixelShuffle的原理是把通道数拆成空间维度，比如输入是CxHxW，输出通道数是C*r^2，然后重排成Cx(rH)x(rW)。这样每个输出像素都由输入像素的加权组合得到，没有重叠区域，自然没有棋盘效应。
# 上采样部分，注意通道数计算upscale=4self.conv_before_upsample=nn.Conv2d(embed_dim,embed_dim*upscale*upscale,3,1,1)self.upsample=nn.PixelShuffle(upscale)self.conv_after_upsample=nn.Conv2d(embed_dim,3,3,1,1)这里有个小技巧：conv_before_upsample的输出通道数必须是embed_dim * upscale * upscale，因为PixelShuffle会把通道数除以r^2。如果upscale=4，通道数变成原来的1/16，所以输入通道要放大16倍。
训练策略：那些容易忽略的细节
损失函数我用L1 + Perceptual Loss的组合。L1保证像素级准确，Perceptual Loss用VGG19的中间层特征，让生成图像在语义上更真实。别用L2，L2会让图像变模糊。
criterion_l1=nn.L1Loss()criterion_perceptual=PerceptualLoss()# 用torchvision的vgg19预训练模型defloss_fn(pred,target):loss_l1=criterion_l1(pred,target)loss_perceptual=criterion_perceptual(pred,target)returnloss_l1+0.1*loss_perceptual学习率调度我用CosineAnnealingWarmRestarts，初始lr=2e-4，T_0=250000，T_mult=2。这个调度器在超分任务上比StepLR效果好，因为模型在训练后期需要更精细的调整。
数据增强方面，我做了随机翻转和旋转90度，但注意不要做颜色抖动，超分任务对颜色一致性要求很高。batch size设16，用梯度累积模拟更大batch，因为SwinIR显存占用大，单卡很难跑大batch。
推理时的trick：重叠切片推理
实际应用时输入图像可能很大（比如1920x1080），直接送进模型会OOM。我采用重叠切片推理：把图像切成256x256的patch，相邻patch重叠16像素，推理完再拼回去。重叠区域用高斯权重融合，避免拼接痕迹。
definference_overlap(model,img,patch_size=256,overlap=16):_,_,H,W=img.shape stride=patch_size-overlap# 计算需要切分的行数和列数h_num=(H-patch_size)//stride+1w_num=(W-patch_size)//stride+1# 创建权重矩阵，中间权重高，边缘权重低weight=torch.ones((1,1,patch_size,patch_size))weight=gaussian_blur(weight,sigma=overlap/2)# 自定义高斯模糊# 逐patch推理并加权融合output=torch.zeros_like(img)weight_sum=torch.zeros_like(img)foriinrange(h_num):forjinrange(w_num):h_start=i*stride w_start=j*stride patch=img[:,:,h_start:h_start+patch_size,w_start:w_start+patch_size]withtorch.no_grad():out_patch=model(patch)output[:,:,h_start:h_start+patch_size,w_start:w_start+patch_size]+=out_patch*weight weight_sum[:,:,h_start:h_start+patch_size,w_start:w_start+patch_size]+=weight output/=weight_sumreturnoutput个人经验总结
窗口大小选择：8x8是平衡计算量和感受野的甜点。窗口太小感受野不足，太大计算量爆炸。如果输入图像纹理丰富（比如人脸），可以试试12x12。
深度配置：官方给的[6,6,6,6]四阶段结构适合通用场景。如果做视频超分，可以加深第一阶段（浅层特征提取），因为视频帧间运动主要在低频区域。
显存优化：训练时用混合精度（torch.cuda.amp），推理时用torch.jit.script加速。我试过把SwinIR转成ONNX，但动态窗口划分导致转换失败，目前还是用PyTorch原生推理。
数据预处理：LR图像用双三次下采样生成，但别用OpenCV的INTER_CUBIC，它和PyTorch的F.interpolate结果有差异。统一用torchvision.transforms.Resize。
调试技巧：第一次跑通后，先拿一张小图（64x64）过一遍，检查输出shape和数值范围。如果输出全是NaN，大概率是softmax前的mask没加对，或者相对位置偏置初始化有问题。
最后说句实在话：SwinIR虽然效果好，但训练成本确实高。如果项目对实时性要求高（比如30fps视频），建议用ESPCN或者FSRCNN，SwinIR更适合离线处理高分辨率图像。但如果你追求极致画质，SwinIR是目前开源方案里最稳的选择。
