Implementing spatio-temporal convolutions in pytorch
import torch as t

class Net_1(t.nn.Module):
    def __init__(self):
        super(Net_1, self).__init__()
        self.conv3d = t.nn.Conv3d(3, 8, (2, 3, 3))
        self.relu = t.nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.conv3d(x))
        return x
   
net_1 = Net_1()

video_1 = t.empty(1, 3, 6, 72, 108).normal_()


%%time
net_1(video_1).shape


Wall time: 18 ms
torch.Size([1, 8, 5, 70, 106])

class Net_2(t.nn.Module):
    def __init__(self):
        super(Net_2, self).__init__()
        self.conv2d = t.nn.Conv2d(3, 8, (3, 3))
        self.conv1d = t.nn.Conv1d(8, 8, (2))
        self.flatten = t.nn.Flatten(start_dim=2)
        self.relu = t.nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.conv2d(x))
        x = self.flatten(x)
        x.transpose_(0, 2)
        x = self.relu(self.conv1d(x))
        x.transpose_(0, 2)
        x.transpose_(0, 1)
        return x.view(5, 8, 70, 106)
        
net_2 = Net_2()

video_2 = t.empty(6, 3, 72, 108).normal_()


%%time
net_2(video_2).shape


Wall time: 58 ms
torch.Size([8, 5, 70, 106])