perf_batch.py

import tensorflow as tf
import tensorflow_addons as tfa
import tensorlayer as tl
import perfplot

tl_model = None
tfa_deform_conv = None
tfa_offset_conv = None


def func(x):
    x_trans = tf.transpose(x, [0, 3, 1, 2])
    res = tfa_deform_conv([x_trans, tfa_offset_conv(x_trans)])
    return tf.transpose(res, [0, 2, 3, 1])


def setup(n):
    global tl_model
    global tfa_deform_conv
    global tfa_offset_conv

    shape = [n, 50, 50, 32]
    x = tf.random.normal(shape)

    ### TensorLayer
    tl_input = tl.layers.Input(shape, name='input')
    tl_offset_conv = tl.layers.Conv2d(n_filter=18, filter_size=(3, 3), strides=(1, 1), padding='SAME', name='offset')(tl_input)
    tl_deform_conv = tl.layers.DeformableConv2d(offset_layer=tl_offset_conv)(tl_input)

    tl_model = tl.models.Model(inputs=tl_input, outputs=tl_deform_conv)
    tl_model.eval()

    tl_model(x)

    ### TensorFlow Addons
    tfa_deform_conv = tfa.layers.DeformableConv2D(filters=32, use_mask=False, padding='same')
    tfa_offset_conv = tf.keras.layers.Conv2D(filters=18, kernel_size=(3, 3), strides=(1, 1), padding='same', data_format='channels_first')

    func(x)

    return x


### perfplot
res = perfplot.bench(
    setup=setup,
    kernels=[
        lambda x: tl_model(x),
        func,
    ],
    labels=["TensorLayer", "TensorFlow Addons"],
    n_range=[1, 2, 4, 8, 16, 32, 64],
    xlabel="Batch size",
    equality_check=None,
)

res.save("perf_batch.png", logx=False, logy=True, transparent=True, bbox_inches="tight")