Extremely slow inference times for CorrelationCost layer on CPU

[evancasey]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
• TensorFlow version and how it was installed (source or binary): 2.0.0, binary
• TensorFlow-Addons version and how it was installed (source or binary): 0.6.0, binary
• Python version: 3.7
• Is GPU used? (yes/no): no, CPU

Describe the bug

Inference takes about ~40s for the example (described below). On my GPU, this runs in under a second.

For reference, the pytorch correlation sampler seems to have much better performance on CPU.

Code to reproduce the issue

import tensorflow as tf
from tensorflow_addons.layers import CorrelationCost

ccl = CorrelationCost(
       kernel_size=1,
       max_displacement=10,
       stride_1=1,
       stride_2=1,
       pad=0,
       data_format='channels_last')

x = tf.ones([1, 148, 276, 1024])
y = tf.ones([1, 148, 276, 1024])

ccl([x,y])

[WindQAQ]

Thanks for the report! Confirmed the issue.
https://colab.research.google.com/drive/1UdYDikjugoFlgunHb5N3arD69R0gHihT

Seems that CPU implementation does not utilize parallelism. Many loops could be computed in parallel here. I will try to improve the performance recently.

[WindQAQ]

Hi @evancasey, the PR has been issued. I test FlowNetC parameters, that is,

ccl = CorrelationCost(
    kernel_size=1,
    max_displacement=10,
    stride_1=1,
    stride_2=2,
    pad=0,
    data_format='channels_last')

x = tf.ones([4, 48, 64, 256])
y = tf.ones([4, 48, 64, 256])

and the runtime is 0.058s (avg over 5 rounds after 5 rounds warm-up). For the parameters in your code snippet, the runtime is about 4s.

I am not so familiar with this area, just copy the FlowNetC parameters from the link you paste above. Please correct me if I do anything wrong. Thanks again for the performance report!

[evancasey]

Incredible. Thank you thank you!!! 🙏 I will give this a shot