深入 Real-ESRGAN 架构：RRDBNet 设计精髓与模型优化全解析

一、引言

上一篇文章我们完成了 Real-ESRGAN 的环境搭建和基础使用。本文将深入剖析其核心生成器RRDBNet的架构设计，并介绍模型导出（ONNX）和推理加速（TensorRT）的完整方案，帮助你在生产环境中高效部署。

二、RRDBNet 架构深度解析

2.1 整体架构

RRDBNet（Residual-in-Residual Dense Block Network）的架构分为三部分：

Input (3×H×W) │ ▼ [Conv 3→64, k=3] ← 浅层特征提取 │ ▼ [RRDB × 23] ← 深层特征提取（核心） │ └─ 每个 RRDB: │ ├─ Dense Block × 3 │ ├─ Residual Scaling (β=0.2) │ └─ Skip Connection │ ▼ [Conv 64→64, k=3] ← 特征融合 │ ▼ [Upsample × 2] ← PixelShuffle 上采样（×4） │ └─ Conv + PixelShuffle(×2) + LReLU │ ▼ [Conv 64→3, k=3] ← 输出重建 │ ▼ Output (3×4H×4W)

2.2 RRDB（Residual-in-Residual Dense Block）

RRDB 是 RRDBNet 的核心构建块，每个 RRDB 内部包含 3 个密集连接的卷积层：

importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassResidualDenseBlock(nn.Module):\"\"\"残差密集块（RDB）\"\"\"def__init__(self,num_feat=64,num_grow_ch=32):super().__init__()self.conv1=nn.Conv2d(num_feat,num_grow_ch,3,1,1)self.conv2=nn.Conv2d(num_feat+num_grow_ch,num_grow_ch,3,1,1)self.conv3=nn.Conv2d(num_feat+2*num_grow_ch,num_grow_ch,3,1,1)self.conv4=nn.Conv2d(num_feat+3*num_grow_ch,num_grow_ch,3,1,1)self.conv5=nn.Conv2d(num_feat+4*num_grow_ch,num_feat,3,1,1)self.lrelu=nn.LeakyReLU(negative_slope=0.2,inplace=True)defforward(self,x):x1=self.lrelu(self.conv1(x))x2=self.lrelu(self.conv2(torch.cat((x,x1),1)))x3=self.lrelu(self.conv3(torch.cat((x,x1,x2),1)))x4=self.lrelu(self.conv4(torch.cat((x,x1,x2,x3),1)))x5=self.conv5(torch.cat((x,x1,x2,x3,x4),1))returnx5*0.2+x# 残差缩放classRRDB(nn.Module):\"\"\"Residual-in-Residual Dense Block\"\"\"def__init__(self,num_feat=64,num_grow_ch=32):super().__init__()self.rdb1=ResidualDenseBlock(num_feat,num_grow_ch)self.rdb2=ResidualDenseBlock(num_feat,num_grow_ch)self.rdb3=ResidualDenseBlock(num_feat,num_grow_ch)defforward(self,x):out=self.rdb1(x)out=self.rdb2(out)out=self.rdb3(out)returnout*0.2+x# 外层残差缩放

2.3 关键设计决策

2.4 参数量与计算量分析

defcount_parameters(model):returnsum(p.numel()forpinmodel.parameters()ifp.requires_grad)defcount_flops(model,input_size=(1,3,256,256)):fromthopimportprofile input_tensor=torch.randn(input_size)flops,params=profile(model,inputs=(input_tensor,))returnflops/1e9,params/1e6# GFLOPs, M paramsmodel=RRDBNet(num_in_ch=3,num_out_ch=3,num_feat=64,num_block=23,num_grow_ch=32,scale=4)print(f"参数:{count_parameters(model)/1e6:.1f}M")# 输出: ~16.7M 参数# 256×256 输入:print(f"计算量:{count_flops(model)[0]:.1f}GFLOPs")# 输出: ~89.2 GFLOPs

三、模型导出与优化

3.1 ONNX 导出

importtorchfrombasicsr.archs.rrdbnet_archimportRRDBNetdefexport_to_onnx():# 加载模型model=RRDBNet(num_in_ch=3,num_out_ch=3,num_feat=64,num_block=23,num_grow_ch=32,scale=4)checkpoint=torch.load('weights/RealESRGAN_x4plus.pth',map_location='cpu')model.load_state_dict(checkpoint['params_ema'],strict=True)model.eval()# 导出 ONNXdummy_input=torch.randn(1,3,256,256)torch.onnx.export(model,dummy_input,'realesrgan_x4plus.onnx',opset_version=17,input_names=['input'],output_names=['output'],dynamic_axes={'input':{0:'batch',2:'height',3:'width'},'output':{0:'batch',2:'height',3:'width'}},do_constant_folding=True)print("ONNX 导出成功: realesrgan_x4plus.onnx")# 验证 ONNX 模型importonnximportonnxruntime onnx_model=onnx.load('realesrgan_x4plus.onnx')onnx.checker.check_model(onnx_model)ort_session=onnxruntime.InferenceSession('realesrgan_x4plus.onnx')ort_input={ort_session.get_inputs()[0].name:dummy_input.numpy()}ort_output=ort_session.run(None,ort_input)print(f"ONNX 推理成功, 输出尺寸:{ort_output[0].shape}")

3.2 TensorRT 加速

# ONNX → TensorRT Enginetrtexec\--onnx=realesrgan_x4plus.onnx\--saveEngine=realesrgan_x4plus.trt\--fp16\--workspace=4096\--optShapes=input:1x3x512x512\--minShapes=input:1x3x64x64\--maxShapes=input:1x3x1024x1024

importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnpclassTRTInference:\"\"\"TensorRT 推理封装\"\"\"def__init__(self,engine_path):self.logger=trt.Logger(trt.Logger.WARNING)withopen(engine_path,'rb')asf:runtime=trt.Runtime(self.logger)self.engine=runtime.deserialize_cuda_engine(f.read())self.context=self.engine.create_execution_context()definfer(self,input_data):# 分配显存d_input=cuda.mem_alloc(input_data.nbytes)d_output=cuda.mem_alloc(self.engine.get_binding_shape(1)[0]*np.prod(self.engine.get_binding_shape(1)[1:])*4)# 推理cuda.memcpy_htod(d_input,input_data)self.context.execute_v2([int(d_input),int(d_output)])output=np.empty(self.engine.get_binding_shape(1),dtype=np.float32)cuda.memcpy_dtoh(output,d_output)returnoutput# 使用trt_infer=TRTInference('realesrgan_x4plus.trt')input_tensor=np.random.randn(1,3,512,512).astype(np.float32)output=trt_infer.infer(input_tensor)

3.3 性能对比

四、生产环境部署架构

┌─────────────────────────────────────────────┐ │ API Gateway │ │ (负载均衡 / 限流 / 认证) │ └─────────────────┬───────────────────────────┘ │ ┌──────────┼──────────┐ ▼ ▼ ▼ ┌─────────┐┌─────────┐┌──────────┐ │ Worker1 ││ Worker2 ││ Worker N │ │ TRT FP16││ TRT FP16││ TRT FP16 │ └─────────┘└─────────┘└──────────┘ │ │ │ └──────────┼──────────┘ ▼ ┌──────────────┐ │ 结果队列 │ │ (Redis/MQ) │ └──────────────┘

五、总结

本文深入剖析了 RRDBNet 的架构设计——残差密集连接、残差缩放、多层 RRDB 堆叠等关键设计如何共同实现高质量的盲超分辨率。同时介绍了 ONNX 导出和 TensorRT 加速的完整流程，TensorRT FP16 推理相比 PyTorch FP32 提速约 4 倍，显存降低约 60%，适合生产环境大规模部署。