PyTorch中DeepLabv3推理后NMS耗时超推理的原因排查问询

我在运行代码时，将输入图像（PyTorch张量形式）送入修改过头部的torchvision.models.segmentation.deeplabv3_mobilenet_v3_large()卷积网络，之后对输出张量执行非极大值抑制（NMS）后处理。但意外发现NMS耗时远超推理过程：

# 执行推理
pred = my_deeplap_based_network(input)

# 创建最大池化操作
R = 9
self.maxpool = torch.nn.MaxPool2d(2*R+1, stride=1, padding=R, ceil_mode=True)

# 运行非极大值抑制，pred是热力图，需要抑制所有非极大值像素
nms = self.detected_points_nms(pred)

def detected_points_nms(self, pred, thresh = 0.3):
        predmax = self.maxpool(pred)
        max_mask = torch.logical_and(pred == predmax,pred > thresh)     
        detected_points = torch.zeros_like(pred)
        detected_points[max_mask] = pred[max_mask]
        return detected_points

输入与输出张量形状为12x3x512x512，在GTX 1080Ti上推理耗时约25ms，NMS耗时约75ms。我原本以为NMS操作应该比DeepLab/MobileNet网络简单更快，这个假设是否错误？还是我的NMS实现存在低效问题？

• 你的假设没错，但当前实现确实存在低效点
  常规目标检测NMS是针对稀疏候选框的操作，但你这里是针对密集热力图的逐像素NMS，计算量本身不小——不过你的实现还有可优化的地方，导致耗时被进一步放大：

1. pred == predmax的逐元素精确比较开销大
   浮点数的逐元素精确相等对比本身效率很低，加上你的张量尺寸是12x3x512x512，整个对比操作会占用大量计算资源。

2. 全零张量创建+索引赋值的冗余操作
   先创建和输入同尺寸的全零张量再做索引赋值，会额外消耗显存带宽，效率远不如直接用掩码与原张量做运算。

3. MaxPool的ceil_mode=True可能引入额外开销
   虽然这里stride=1不会改变输出尺寸，但ceil_mode=True会增加边缘区域的池化计算逻辑，带来不必要的性能损耗。

• 优化后的NMS实现
  可以通过以下修改降低耗时：

  def detected_points_nms(self, pred, thresh=0.3):
      predmax = self.maxpool(pred)
      # 用近似相等替代精确相等，避免浮点数精确对比的高开销
      max_mask = torch.logical_and(torch.isclose(pred, predmax, atol=1e-6), pred > thresh)
      # 直接用掩码乘原张量，替代全零张量+赋值的冗余操作
      detected_points = pred * max_mask.float()
      return detected_points
  

• 进一步优化方向

  • 如果任务允许，缩小池化核尺寸R，或者先将热力图下采样再做NMS、最后上采样回原尺寸，能大幅减少计算量。
  • 将MaxPool和NMS操作融合到模型前向传播中，使用PyTorch的JIT编译器（torch.jit.trace/torch.jit.script）做算子融合优化，提升整体效率。
  • 开启PyTorch的CUDA优化（如torch.backends.cudnn.benchmark = True），确保算子使用最优的CUDA实现。

内容的提问来源于stack exchange，提问作者Philipp