PyTorch实现多Atrous Convolution层后训练变慢GPU利用率低如何解决

性能瓶颈分析 #

• 大dilation空洞卷积的访存效率低下：dilation=6、dilation=9的3x3卷积采样的像素在显存中不连续，GPU缓存命中率大幅降低，大部分时间卡在等待显存数据返回，直接导致GPU计算利用率暴跌。
• 多独立算子顺序执行浪费算力：三个空洞卷积、全局池化分支互相之间没有数据依赖，但PyTorch默认单CUDA流下会顺序执行所有算子，GPU计算单元存在大量闲置。
• 冗余的上采样计算：全局平均池化输出为1x1尺寸的特征，使用双线性插值上采样到16x16属于完全冗余的计算，所有输出点值完全相同，额外占用了算力和显存带宽。
• 无前置通道降维带来的额外开销：标准ASPP结构会先通过1x1卷积降低输入通道数再计算多支路空洞卷积，你的实现直接在1024通道输入上计算3个512输出的空洞卷积，计算量和访存量是最优实现的4倍以上。

优化方案 #

• 增加前置通道降维：在空洞卷积分支前新增1个1x1卷积，将1024通道输入降到256通道后再分别输入三个空洞卷积，可直接降低75%的计算量和访存开销。
• 替换上采样实现：将nn.Upsample替换为张量扩展操作，修改为self.relu(self.conv1x1(self.gap(x))).expand(-1, -1, 16, 16)，等价输出下完全消除插值计算开销。
• 开启cuDNN基准模式：在训练代码开头添加torch.backends.cudnn.benchmark = True，让cuDNN自动搜索当前尺寸、dilation参数下的最优卷积Kernel，可提升大dilation卷积30%以上的执行速度。
• 启用多流并行：将三个独立的空洞卷积发射到不同的CUDA流执行，消除顺序执行的等待开销，进一步提升GPU利用率。
• 减少临时变量：直接在torch.cat中传入四个分支的输出，无需用list存储临时特征，降低显存分配和拷贝开销。

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