深入 Real-ESRGAN 架构RRDBNet 设计精髓与模型优化全解析一、引言上一篇文章我们完成了 Real-ESRGAN 的环境搭建和基础使用。本文将深入剖析其核心生成器RRDBNet的架构设计并介绍模型导出ONNX和推理加速TensorRT的完整方案帮助你在生产环境中高效部署。二、RRDBNet 架构深度解析2.1 整体架构RRDBNetResidual-in-Residual Dense Block Network的架构分为三部分Input (3×H×W) │ ▼ [Conv 3→64, k3] ← 浅层特征提取 │ ▼ [RRDB × 23] ← 深层特征提取核心 │ └─ 每个 RRDB: │ ├─ Dense Block × 3 │ ├─ Residual Scaling (β0.2) │ └─ Skip Connection │ ▼ [Conv 64→64, k3] ← 特征融合 │ ▼ [Upsample × 2] ← PixelShuffle 上采样×4 │ └─ Conv PixelShuffle(×2) LReLU │ ▼ [Conv 64→3, k3] ← 输出重建 │ ▼ Output (3×4H×4W)2.2 RRDBResidual-in-Residual Dense BlockRRDB 是 RRDBNet 的核心构建块每个 RRDB 内部包含 3 个密集连接的卷积层importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassResidualDenseBlock(nn.Module):\\\残差密集块RDB\\\def__init__(self,num_feat64,num_grow_ch32):super().__init__()self.conv1nn.Conv2d(num_feat,num_grow_ch,3,1,1)self.conv2nn.Conv2d(num_featnum_grow_ch,num_grow_ch,3,1,1)self.conv3nn.Conv2d(num_feat2*num_grow_ch,num_grow_ch,3,1,1)self.conv4nn.Conv2d(num_feat3*num_grow_ch,num_grow_ch,3,1,1)self.conv5nn.Conv2d(num_feat4*num_grow_ch,num_feat,3,1,1)self.lrelunn.LeakyReLU(negative_slope0.2,inplaceTrue)defforward(self,x):x1self.lrelu(self.conv1(x))x2self.lrelu(self.conv2(torch.cat((x,x1),1)))x3self.lrelu(self.conv3(torch.cat((x,x1,x2),1)))x4self.lrelu(self.conv4(torch.cat((x,x1,x2,x3),1)))x5self.conv5(torch.cat((x,x1,x2,x3,x4),1))returnx5*0.2x# 残差缩放classRRDB(nn.Module):\\\Residual-in-Residual Dense Block\\\def__init__(self,num_feat64,num_grow_ch32):super().__init__()self.rdb1ResidualDenseBlock(num_feat,num_grow_ch)self.rdb2ResidualDenseBlock(num_feat,num_grow_ch)self.rdb3ResidualDenseBlock(num_feat,num_grow_ch)defforward(self,x):outself.rdb1(x)outself.rdb2(out)outself.rdb3(out)returnout*0.2x# 外层残差缩放2.3 关键设计决策设计要素值原因特征通道数num_feat64平衡性能与参数量扩大效果提升有限生长通道num_grow_ch32密集连接的增长通道数控制计算量RRDB 数量23深层网络充足的感受野残差缩放 β0.2防止梯度爆炸稳定训练密集连接数每 RDB 5 层最大化信息流动LeakyReLU 斜率0.2保留负值信息流避免 Dead ReLU2.4 参数量与计算量分析defcount_parameters(model):returnsum(p.numel()forpinmodel.parameters()ifp.requires_grad)defcount_flops(model,input_size(1,3,256,256)):fromthopimportprofile input_tensortorch.randn(input_size)flops,paramsprofile(model,inputs(input_tensor,))returnflops/1e9,params/1e6# GFLOPs, M paramsmodelRRDBNet(num_in_ch3,num_out_ch3,num_feat64,num_block23,num_grow_ch32,scale4)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():# 加载模型modelRRDBNet(num_in_ch3,num_out_ch3,num_feat64,num_block23,num_grow_ch32,scale4)checkpointtorch.load(weights/RealESRGAN_x4plus.pth,map_locationcpu)model.load_state_dict(checkpoint[params_ema],strictTrue)model.eval()# 导出 ONNXdummy_inputtorch.randn(1,3,256,256)torch.onnx.export(model,dummy_input,realesrgan_x4plus.onnx,opset_version17,input_names[input],output_names[output],dynamic_axes{input:{0:batch,2:height,3:width},output:{0:batch,2:height,3:width}},do_constant_foldingTrue)print(ONNX 导出成功: realesrgan_x4plus.onnx)# 验证 ONNX 模型importonnximportonnxruntime onnx_modelonnx.load(realesrgan_x4plus.onnx)onnx.checker.check_model(onnx_model)ort_sessiononnxruntime.InferenceSession(realesrgan_x4plus.onnx)ort_input{ort_session.get_inputs()[0].name:dummy_input.numpy()}ort_outputort_session.run(None,ort_input)print(fONNX 推理成功, 输出尺寸:{ort_output[0].shape})3.2 TensorRT 加速# ONNX → TensorRT Enginetrtexec\--onnxrealesrgan_x4plus.onnx\--saveEnginerealesrgan_x4plus.trt\--fp16\--workspace4096\--optShapesinput:1x3x512x512\--minShapesinput:1x3x64x64\--maxShapesinput:1x3x1024x1024importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnpclassTRTInference:\\\TensorRT 推理封装\\\def__init__(self,engine_path):self.loggertrt.Logger(trt.Logger.WARNING)withopen(engine_path,rb)asf:runtimetrt.Runtime(self.logger)self.engineruntime.deserialize_cuda_engine(f.read())self.contextself.engine.create_execution_context()definfer(self,input_data):# 分配显存d_inputcuda.mem_alloc(input_data.nbytes)d_outputcuda.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)])outputnp.empty(self.engine.get_binding_shape(1),dtypenp.float32)cuda.memcpy_dtoh(output,d_output)returnoutput# 使用trt_inferTRTInference(realesrgan_x4plus.trt)input_tensornp.random.randn(1,3,512,512).astype(np.float32)outputtrt_infer.infer(input_tensor)3.3 性能对比推理方式512×512 推理时间4096×4096 推理时间显存占用PyTorch (FP32)180ms5400ms6.2GBPyTorch (FP16)95ms3100ms3.8GBONNX Runtime140ms4800ms5.1GBTensorRT (FP16)45ms1500ms2.4GBTensorRT (INT8)28ms980ms1.6GB四、生产环境部署架构┌─────────────────────────────────────────────┐ │ API Gateway │ │ (负载均衡 / 限流 / 认证) │ └─────────────────┬───────────────────────────┘ │ ┌──────────┼──────────┐ ▼ ▼ ▼ ┌─────────┐┌─────────┐┌──────────┐ │ Worker1 ││ Worker2 ││ Worker N │ │ TRT FP16││ TRT FP16││ TRT FP16 │ └─────────┘└─────────┘└──────────┘ │ │ │ └──────────┼──────────┘ ▼ ┌──────────────┐ │ 结果队列 │ │ (Redis/MQ) │ └──────────────┘五、总结本文深入剖析了 RRDBNet 的架构设计——残差密集连接、残差缩放、多层 RRDB 堆叠等关键设计如何共同实现高质量的盲超分辨率。同时介绍了 ONNX 导出和 TensorRT 加速的完整流程TensorRT FP16 推理相比 PyTorch FP32 提速约 4 倍显存降低约 60%适合生产环境大规模部署。