在pytorch中使用可分离的二维卷积实现三维高斯模糊

我试图在pytorch中实现一个类似高斯的3D体积模糊。我可以很容易地用二维高斯核卷积来对二维图像进行二维模糊，同样的方法似乎也适用于三维高斯核。但是，它在3D中的速度非常慢（特别是对于较大的sigmas/内核大小）。我知道这也可以通过使用2D内核进行3次卷积来实现，这应该要快得多，但我无法让它正常工作。我的测试用例如下

import torch
import torch.nn.functional as F

VOL_SIZE = 21


def make_gaussian_kernel(sigma):
    ks = int(sigma * 5)
    if ks % 2 == 0:
        ks += 1
    ts = torch.linspace(-ks // 2, ks // 2 + 1, ks)
    gauss = torch.exp((-(ts / sigma)**2 / 2))
    kernel = gauss / gauss.sum()

    return kernel


def test_3d_gaussian_blur(blur_sigma=2):
    # Make a test volume
    vol = torch.zeros([VOL_SIZE] * 3)
    vol[VOL_SIZE // 2, VOL_SIZE // 2, VOL_SIZE // 2] = 1

    # 3D convolution
    vol_in = vol.reshape(1, 1, *vol.shape)
    k = make_gaussian_kernel(blur_sigma)
    k3d = torch.einsum('i,j,k->ijk', k, k, k)
    k3d = k3d / k3d.sum()
    vol_3d = F.conv3d(vol_in, k3d.reshape(1, 1, *k3d.shape), stride=1, padding=len(k) // 2)

    # Separable 2D convolution
    vol_in = vol.reshape(1, *vol.shape)
    k2d = torch.einsum('i,j->ij', k, k)
    k2d = k2d / k2d.sum()
    k2d = k2d.expand(VOL_SIZE, 1, *k2d.shape)
    for i in range(3):
        vol_in = vol_in.permute(0, 3, 1, 2)
        vol_in = F.conv2d(vol_in, k2d, stride=1, padding=len(k) // 2, groups=VOL_SIZE)
    vol_3d_sep = vol_in

    torch.allclose(vol_3d, vol_3d_sep)  # --> False

---

任何帮助都将不胜感激

从理论上讲，可以使用三个2d卷积计算3d高斯卷积，但这意味着您必须减小2d内核的大小，因为在每个方向上有效地卷积两次

但在计算上更有效的方法（也是你通常想要的）是将核分离成一维核。我更改了函数的第二部分来实现这一点。（我必须说我真的很喜欢你的基于排列的方法！）因为你使用的是3d卷，你不能很好地使用

conv2d

或

conv1d

函数，所以最好的方法就是使用

conv3d

，即使你只是计算1d卷积

请注意，

allclose

使用的阈值为

1e-8

，我们使用此方法无法达到该阈值，可能是由于取消错误

def测试\u 3d\u高斯\u模糊（模糊\u西格玛=2）：
#做一个测试卷
vol=torch.randn（[vol_SIZE]*3）#使用非零的值
卷[vol\u SIZE//2，vol\u SIZE//2，vol\u SIZE//2]=1
#三维卷积
体积=体积重塑（1，1，*体积形状）
k=使之成为高斯核（模糊西格玛）
k3d=torch.einsum（'i，j，k->ijk'，k，k，k）
k3d=k3d/k3d.sum（）
vol_3d=F.conv3d（vol_in，k3d.revorme（1,1，*k3d.shape），步长=1，填充=len（k）//2）
#可分一维卷积
vol_in=vol[无，无，…]
#k2d=火炬。einsum（'i，j->ij'，k，k）
#k2d=k2d/k2d.sum（）#如果内核的和已经为零，则不需要检查：
#打印（f'{k2d.sum（）=}'）
k1d=k[无，无，：，无，无]
对于范围（3）中的i：
vol_in=vol_in.排列（0,1,4,2,3）
vol_-in=F.conv3d（vol_-in，k1d，步长=1，填充=（len（k）//2,0,0））
vol_3d_sep=vol_in
打印（（vol_3d-vol_3d_sep）.abs（）.max（））#某物~1e-7
打印（torch.ALLCOSE（第3卷，第3卷，第9页））#ALLCOSE检查是否在1e-8左右

---

附录：如果您确实想滥用

conv2d

来处理卷，您可以尝试

# separate 3d kernel into 1d + 2d
vol_in = vol[None, None, ...]
k2d = torch.einsum('i,j->ij', k, k)
k2d = k2d.expand(VOL_SIZE, 1, len(k), len(k))
# k2d = k2d / k2d.sum() # not necessary if kernel already sums to zero, check:
# print(f'{k2d.sum()=}')
k1d = k[None, None, :, None, None]
vol_in = F.conv3d(vol_in, k1d, stride=1, padding=(len(k) // 2, 0, 0))
vol_in = vol_in[0, ...]
# abuse conv2d-groups argument for volume dimension, works only for 1 channel volumes
vol_in = F.conv2d(vol_in, k2d, stride=1, padding=(len(k) // 2, len(k) // 2), groups=VOL_SIZE)
vol_3d_sep = vol_in

或者专门使用

conv2d

可以执行以下操作：

# separate 3d kernel into 1d + 2d
vol_in = vol[None,  ...]
# 1d kernel
k1d = k[None, None, :,  None]
k1d = k1d.expand(VOL_SIZE, 1, len(k), 1)
# 2d kernel
k2d = torch.einsum('i,j->ij', k, k)
k2d = k2d.expand(VOL_SIZE, 1, len(k), len(k))
vol_in = vol_in.permute(0, 2, 1, 3)
vol_in = F.conv2d(vol_in, k1d, stride=1, padding=(len(k) // 2, 0), groups=VOL_SIZE)
vol_in = vol_in.permute(0, 2, 1, 3)
vol_in = F.conv2d(vol_in, k2d, stride=1, padding=(len(k) // 2, len(k) // 2), groups=VOL_SIZE)
vol_3d_sep = vol_in

---

这些应该仍然比三个连续的二维卷积快。

实际上，你可以将一个三维（各向同性）高斯核分离为三个一维核。（在理论上，分离成2d内核是可行的，但没有理由这么做！）是的，我也这么想，但在2d的情况下，我可以将第三维视为通道，并使用分组2d卷积（至少我正在尝试）。在1d的情况下，我有两个空间维度，所以不知道如何在pytorch中正确实现它。可能会进行一些重塑？在任何情况下，您都可以使用通常的conv3d功能，但使用

k x 1 x 1

内核。看起来很棒，谢谢！它甚至使用零+1:）来传递断言。在我的实际问题上，这个版本与3d conv内核方法相比，速度提高了20-25%。尽管考虑到gpu上二维卷积的速度，我希望看到比这多一点。我仍然想让2d conv版本进行全面比较，但这肯定是一个很大的改进。@不平衡您有任何理由相信2d卷积会更快吗？我希望它们（至少在cpu上需要（k的长度）多几倍）。通常，根据硬件和内核的大小，具有一个3d卷积的内核可能会更快（或者相反，对于非常大的内核，3x1d方法肯定会快得多）。你用更大的内核/输入测试过它吗？我的意思是你能做的只是一个二维卷积，然后是一个一维卷积，但在我的书中，三个二维卷积似乎没有多大意义。或者你想使用三个二维卷积有什么特别的原因吗？当然，我同意一维卷积应该是最有效的方法。我没有任何基准测试，但处理成批的2d图像（~64*200^2）需要经过几十次卷积，但这里使用3个“1d”卷积的1x100^3似乎需要更长的时间。虽然blur（~11）的内核大小比CNN（~3x3）的要大，但我还是觉得还有很多需要改进的地方。我可能错了，或者在其他地方浪费了时间……内核大小对性能有很大的影响！将

nxn

映像与

kxk

内核进行卷积是一个

O（n^2*k^2）

操作。可以尝试加快计算速度的一件事是设置

torch.backends.cudnn.benchmark=True

。我添加了两个片段，它们“滥用”

conv2d

在一个卷上执行模糊处理，我很确定它们比使用三个二维卷积要好。但您必须记住，这些仅适用于单通道图像。