Question

我有一些形状为 B*C*T*H*W 的数据。我想在 H*W 维度上应用 2d 卷积。

有两个选项（我看到的）：

应用形状为 (1, 3, 3) 的部分 3D 卷积。 3D 卷积接受形状为 B*C*T*H*W 的数据，这正是我所拥有的。然而，这是一个可能优化不足的伪 3d 转换（尽管它在 P3D 网络中被大量使用）。
转置数据，应用 2D 卷积，并将结果转回。这需要数据整形的开销，但它利用了高度优化的 2D 卷积。

data = raw_data.transpose(1,2).reshape(b*t, c, h, w).detach()
out = conv2d(data)
out = out.view(b, t, c, h, w).transpose(1, 2).contiguous()

哪个更快？

（注意：我在下面有一个自我回答。这旨在为那些使用谷歌搜索的人（也就是 20 分钟前的我）做一个快速笔记）

Answer 1

环境：PyTorch 1.7.1、CUDA 11.0、RTX 2080 TI。

TL;DR：转置 + 2D 转换速度更快（在这种环境下，对于测试的数据形状）。

代码（修改自 here）：

import torch
import torch.nn as nn
import time

b = 4
c = 64
t = 4
h = 256
w = 256
raw_data = torch.randn(b, c, t, h, w).cuda()

def time2D():
        conv2d = nn.Conv2d(c, c, kernel_size=3, padding=1).cuda()
        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            data = raw_data.transpose(1,2).reshape(b*t, c, h, w).detach()
            out = conv2d(data)
            out = out.view(b, t, c, h, w).transpose(1, 2).contiguous()
            out.mean().backward()
        torch.cuda.synchronize()
        end = time.time()
        print("  --- %s ---  " %(end - start))

def time3D():
        conv3d = nn.Conv3d(c, c, kernel_size=(1,3,3), padding=(0,1,1)).cuda()
        torch.cuda.synchronize()
        start = time.time()
        for _ in range(100):
            out = conv3d(raw_data.detach())
            out.mean().backward()
        torch.cuda.synchronize()
        end = time.time()
        print("  --- %s ---  " %(end - start))

print("Initializing cuda state")
time2D()

print("going to time2D")
time2D()
print("going to time3D")
time3D()

对于形状 = 4*64*4*256*256：

2D: 1.8675172328948975
3D: 4.384545087814331

对于形状 = 8*512*16*64*64：

2D: 37.95961904525757
3D: 49.730860471725464

对于形状 = 4*128*128*16*16：

2D: 0.6455907821655273
3D: 1.8380646705627441

部分3D卷积还是转置+2D卷积更快？

1 个答案: