深入 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 个密集连接的卷积层：
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class ResidualDenseBlock(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)
def forward(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))
return x5 * 0.2 + x # 残差缩放
class RRDB(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)
def forward(self, x):
out = self.rdb1(x)
out = self.rdb2(out)
out = self.rdb3(out)
return out * 0.2 + x # 外层残差缩放2.3 关键设计决策
| 设计要素 | 值 | 原因 |
|---|---|---|
特征通道数 num_feat |
64 | 平衡性能与参数量，扩大效果提升有限 |
生长通道 num_grow_ch |
32 | 密集连接的增长通道数，控制计算量 |
| RRDB 数量 | 23 | 深层网络，充足的感受野 |
| 残差缩放 β | 0.2 | 防止梯度爆炸，稳定训练 |
| 密集连接数 | 每 RDB 5 层 | 最大化信息流动 |
| LeakyReLU 斜率 | 0.2 | 保留负值信息流，避免 Dead ReLU |
2.4 参数量与计算量分析
python
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def count_flops(model, input_size=(1, 3, 256, 256)):
from thop import profile
input_tensor = torch.randn(input_size)
flops, params = profile(model, inputs=(input_tensor,))
return flops / 1e9, params / 1e6 # GFLOPs, M params
model = 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 导出
python
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
def export_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()
# 导出 ONNX
dummy_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 模型
import onnx
import onnxruntime
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 加速
bash
# ONNX → TensorRT Engine
trtexec \
--onnx=realesrgan_x4plus.onnx \
--saveEngine=realesrgan_x4plus.trt \
--fp16 \
--workspace=4096 \
--optShapes=input:1x3x512x512 \
--minShapes=input:1x3x64x64 \
--maxShapes=input:1x3x1024x1024
python
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
class TRTInference:
\"\"\"TensorRT 推理封装\"\"\"
def __init__(self, engine_path):
self.logger = trt.Logger(trt.Logger.WARNING)
with open(engine_path, 'rb') as f:
runtime = trt.Runtime(self.logger)
self.engine = runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
def infer(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)
return output
# 使用
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 性能对比
| 推理方式 | 512×512 推理时间 | 4096×4096 推理时间 | 显存占用 |
|---|---|---|---|
| PyTorch (FP32) | 180ms | 5400ms | 6.2GB |
| PyTorch (FP16) | 95ms | 3100ms | 3.8GB |
| ONNX Runtime | 140ms | 4800ms | 5.1GB |
| TensorRT (FP16) | 45ms | 1500ms | 2.4GB |
| TensorRT (INT8) | 28ms | 980ms | 1.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%，适合生产环境大规模部署。