Gaussian agent 46% slower in Pytorch compared to Tensorflow

Rick: <p>We are currently trying to transfer our department code base from Tensorflow to Pytorch. For most of our code, this transformation has been done. In doing this we greatly simplified our codebase while not having to sacrifice on performance <img src="https://discuss.pytorch.org/images/emoji/apple/rocket.png?v=9" title=":rocket:" class="emoji" alt=":rocket:"> .</p>
<p>For one agent, however, we noticed that the Pytorch Version is about 46% slower than it’s Tensorflow counterpart (while using the <a href="https://www.tensorflow.org/api_docs/python/tf/function" rel="noopener nofollow ugc">tf.function decorator</a>). As expected when the <a href="https://www.tensorflow.org/api_docs/python/tf/function" rel="noopener nofollow ugc">tf.function</a> decorator is not used Pytorch becomes about 10 times faster than TensorFlow. I profiled both versions using the <a href="https://functiontrace.com/" rel="noopener nofollow ugc">functiontrace</a> and it looks like the biggest difference can be found in the forward and backward passes through the network. This difference is most prevalent in the Actor-network. This network contains the following components:</p>
<ul>
<li>A pass through a fully connected network.</li>
<li>A <a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.cauchy.Cauchy.rsample" rel="noopener nofollow ugc">rsample operation</a> from a normal distribution (Explicitly performed in tensorflow).</li>
<li>A operation that calculates the <a href="https://pytorch.org/docs/stable/_modules/torch/distributions/continuous_bernoulli.html#ContinuousBernoulli.log_prob" rel="noopener nofollow ugc">log probabilities</a> of these sampled actions.</li>
<li>A operation which performs a squashing operation on the normal distribution.</li>
</ul>
<p>A timed comparison of each of these components between Pytorch and Tensorflow can be found in this issue](<a href="https://github.com/rickstaa/torch-tf2-lac-speed-compare/issues/1" rel="noopener nofollow ugc">https://github.com/rickstaa/torch-tf2-lac-speed-compare/issues/1</a>). The code used for these tests can be found in <a href="https://github.com/rickstaa/torch-tf2-lac-speed-compare" rel="noopener nofollow ugc">the accompanying repository</a>.</p>
<p>The conclusion from these tests is as follows:</p>
<ul>
<li>The sample operation is significantly faster in Pytorch.</li>
<li>The log probability operation has the same performance in both versions.</li>
<li>The squashing operation is significantly slower in Pytorch.</li>
</ul>
<p>We can therefore see that the difference is caused by the Squahsing operations. This is also the part which differs the most between the two versions. As to my knowledge, no bijector class exists in Pytorch the squashing operation has to be performed explicitly (see <a href="https://github.com/rickstaa/torch-tf2-lac-speed-compare/blob/c1c43f15cd58999bd63f6073dd4aa89a3d93cdb4/gaussian_actor_torch.py#L108-L117" rel="noopener nofollow ugc">L108-L117 of gaussian actor_torch</a>):</p>
<pre><code class="lang-auto">logp_pi = pi_distribution.log_prob(pi_action).sum(axis=-1)
logp_pi -= (2 * (np.log(2) - pi_action - F.softplus(-2 * pi_action))).sum(
    axis=sum_axis
)
</code></pre>
<p>The Tensorflow version, however, uses <a href="https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Bijector" rel="noopener nofollow ugc">the TensorFlow bijection class</a> to perform the squashing operation. This bijection is defined using the <a href="https://github.com/rickstaa/torch-tf2-lac-speed-compare/blob/main/squash_bijector.py#L9" rel="noopener nofollow ugc">squashed_bijector class</a>):</p>
<pre><code class="lang-auto">"""Creates a squash (tanh) bijector for using the parameterization trick.
"""

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp


class SquashBijector(tfp.bijectors.Bijector):
    """A squash bijector used to keeps track of the distribution properties when the
    distribution is transformed using the tanh squash function."""

    def __init__(self, validate_args=False, name="tanh"):
        super(SquashBijector, self).__init__(
            forward_min_event_ndims=0, validate_args=validate_args, name=name
        )

    def _forward(self, x):
        return tf.nn.tanh(x)
        # return x

    def _inverse(self, y):
        return tf.atanh(y)

    def _forward_log_det_jacobian(self, x):
        return 2.0 * (
            np.log(2.0) - x - tf.nn.softplus(-2.0 * x)
        )  # IMPROVE: Speed check
</code></pre>
<p>This bijector is then used in <a href="https://github.com/rickstaa/torch-tf2-lac-speed-compare/blob/c1c43f15cd58999bd63f6073dd4aa89a3d93cdb4/gaussian_actor_tf2.py#L121-L137" rel="noopener nofollow ugc">L121-L137 of the gaussian actor</a> to perform the squashing operation:</p>
<pre><code class="lang-auto"># Create bijectors (Used in the re-parameterization trick)
squash_bijector = SquashBijector()
affine_bijector = tfp.bijectors.Shift(mu)(tfp.bijectors.Scale(sigma))

# Transform distribution back to the original policy distribution
reparm_trick_bijector = tfp.bijectors.Chain((squash_bijector, affine_bijector))
distribution = tfp.distributions.TransformedDistribution(
    distribution=base_distribution, bijector=reparm_trick_bijector
)
</code></pre>
<p>I, therefore, think that it is <a href="https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Bijector" rel="noopener nofollow ugc">the bijector class</a> in combination with the <a href="https://www.tensorflow.org/api_docs/python/tf/function" rel="noopener nofollow ugc">tf.function wrapper</a> that allows TensorFlow to be able to perform the squashing operation in a more efficient way. I tried to use the TorchSCript wrapper to speed up this computation. This however currently is not yet supported as torch script wrapper does not support the normal distributions that are used in the gaussian actor (see pytorch/pytorch#18094 and pytorch/pytorch#29843).</p>
<p>Since I am quite new to TensorFlow I was wondering if a more experienced user can comment on my reasoning and see if I missed something. Thanks a lot in advance!</p>
<p><strong>Additional Information:</strong></p>
<p>The full code base of the PyTorch and TF2 implementations of the agent can be found in <a href="https://github.com/rickstaa/LAC-TF2-TORCH-translation" rel="noopener nofollow ugc">this repository</a>. I’m using the following system:</p>
<ul>
<li>
<strong>PC:</strong> Hp Zbook x360</li>
<li>
<strong>OS:</strong> Ubuntu 20.04</li>
<li>
<strong>CPU:</strong> Intel® Core™ i7-8750H CPU @ 2.20GHz</li>
<li>
<strong>GPU:</strong> Quadro P1000</li>
<li>
<strong>Torch version:</strong>  1.7.1</li>
<li>
<strong>TensorFlow version:</strong> 2.3.0</li>
<li>
<strong>Tensorflow-distributions version:</strong>  0.11.0</li>
</ul>