如何对3D体图像沿三个正交轴执行2D convolution替代3D卷积降低算力开销

该需求完全有成熟可行的实现路径，这一思路也是目前工业界和学术界降低3D卷积算力开销的主流方案之一，被称为正交2D卷积或伪3D卷积，实际落地中可以在精度损失可控的前提下将3D卷积的算力开销降低70%以上。

核心实现逻辑 #

常规3D体数据的维度为(N, C, D, H, W)，其中N为批量大小、C为通道数、D为深度、H为高度、W为宽度，三个方向的2D卷积分别按如下逻辑处理：

• xy平面方向：将深度D作为临时批量维度，把输入reshape为(N*D, C, H, W)后执行普通2D卷积，处理完成后还原维度为(N, C_out, D, H, W)，得到xy方向的特征体
• xz平面方向：调整维度顺序将H移到临时批量维度，把输入reshape为(N*H, C, D, W)后执行普通2D卷积，处理完成后调整维度顺序还原为(N, C_out, D, H, W)，得到xz方向的特征体
• yz平面方向：调整维度顺序将W移到临时批量维度，把输入reshape为(N*W, C, D, H)后执行普通2D卷积，处理完成后调整维度顺序还原为(N, C_out, D, H, W)，得到yz方向的特征体

多方向特征融合方案 #

得到三个方向的特征体后，可以根据业务需求选择不同的融合方式：

• 通道拼接：将三个特征体在通道维度拼接，总通道数为3*C_out，后续接入普通卷积做特征融合，精度最高，适合对效果要求高的场景
• 加权相加：给三个方向的特征分配可学习的权重，加权求和得到单个体特征，算力开销更低，适合轻量化场景
• 逐元素相加：直接将三个特征体逐元素相加，计算量最小，适合边缘端部署场景

PyTorch实现示例 #

import torch
import torch.nn as nn

class Ortho2DConv3D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, padding=1):
        super().__init__()
        # 三个方向的2D卷积，参数可共享也可独立，默认独立参数效果更好
        self.conv_xy = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
        self.conv_xz = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
        self.conv_yz = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
    
    def forward(self, x):
        N, C, D, H, W = x.shape
        # xy方向卷积
        x_xy = x.permute(0,2,1,3,4).reshape(N*D, C, H, W)
        out_xy = self.conv_xy(x_xy).reshape(N, D, -1, H, W).permute(0,2,1,3,4)
        # xz方向卷积
        x_xz = x.permute(0,3,1,2,4).reshape(N*H, C, D, W)
        out_xz = self.conv_xz(x_xz).reshape(N, H, -1, D, W).permute(0,2,3,1,4)
        # yz方向卷积
        x_yz = x.permute(0,4,1,2,3).reshape(N*W, C, D, H)
        out_yz = self.conv_yz(x_yz).reshape(N, W, -1, D, H).permute(0,2,3,4,1)
        # 此处默认使用通道拼接，可自行修改为相加、加权相加等融合逻辑
        return torch.cat([out_xy, out_xz, out_yz], dim=1)

该模块可以直接替换原生3D卷积，嵌入到任意3D任务网络（如医学影像分割、三维点云体素分类等）中使用。

注意事项 #

• 三个方向的2D卷积的kernel_size、padding参数建议设置为一致，保证输出空间尺寸匹配，避免融合时维度不兼容
• 针对各向异性的3D体数据，可以适当调高采样率更低的方向的卷积权重，降低精度损失
• 如果需要进一步压缩算力，可以让三个方向的2D卷积共享权重，能额外降低2/3的参数量

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