Compare 1D Separable Convolutions in PyTorch and Tensorflow

benchmark_separable_convs.py

#!/usr/bin/env python

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import sys
import torch
import numpy as np
import tensorflow as tf
from time import perf_counter
from torch.nn import Module, Conv1d
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import SeparableConv1D

tf.compat.v1.disable_eager_execution()


class TCSConv1d(Module):
    def __init__(self, filters, kernel_size):
        super(TCSConv1d, self).__init__()
        self.depthwise = Conv1d(filters, filters, kernel_size=kernel_size, groups=filters)
        self.pointwise = Conv1d(filters, filters, kernel_size=1)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x


def pytorch_model(N, C, W, K):
    model = TCSConv1d(C, K)
    model.to('cuda')
    data = torch.ones((N, C, W), dtype=torch.float32)
    return model, data


def tensorflow_model(N, C, W, K):
    inputs = Input(batch_shape=(N, C, W))
    x = SeparableConv1D(C, K, data_format='channels_first')(inputs)
    model = Model(inputs=inputs, outputs=x)
    model.compile(loss='mean_squared_error', optimizer='adam')
    data = np.ones((N, C, W), dtype=np.float32)
    return model, data


def benchmark(N=100, W=2000):

    test_sizes = [
        (256, 33),
        (256, 39),
        (512, 51),
        (512, 63),
        (512, 75),
        (512, 86),
    ]

    for C, K in test_sizes:

        model, data = tensorflow_model(N, C, W, K)
        model.predict(data)
        t0 = perf_counter()
        model.predict(data)
        tftime = perf_counter() - t0

        model, data = pytorch_model(N, C, W, K)
        t0 = perf_counter()
        model(data.to('cuda')).to('cpu')
        torch.cuda.synchronize('cuda')
        pttime = perf_counter() - t0

        print("Filters:", C, "Kernel Size:", K)
        print("PyTorch    %.4f" % pttime)
        print("Tensorflow %.4f" % tftime)


if __name__ == "__main__":
    benchmark()

On a GV100 comparing PyTorch 1.3.1 and Tensorflow 2.0 (cuda 10.1), I make PyTorch up to twice as fast in some cases.

Filters: 256 Kernel Size: 33
PyTorch    0.2585
Tensorflow 0.3271
Filters: 256 Kernel Size: 39
PyTorch    0.1684
Tensorflow 0.3174
Filters: 512 Kernel Size: 51
PyTorch    0.3588
Tensorflow 0.5791
Filters: 512 Kernel Size: 63
PyTorch    0.3329
Tensorflow 0.5758
Filters: 512 Kernel Size: 75
PyTorch    0.3353
Tensorflow 0.6062
Filters: 512 Kernel Size: 86
PyTorch    0.3292
Tensorflow 0.6070