PyTorch slower than Tensorflow/Keras during training

ttiago_jm: <p>Hi,<br>
I’m porting a Tensorflow project to PyTorch and after executing I realize that Tensorflow is faster than PyTorch during training and it’s weird to me.</p>
<p>I isolate the 2 projects in 2 different notebooks with all classes and functions needed to train the model.<br>
Remember this is just a snippet of a bigger project (a project that uses genetic algorithms) and that’s why CustomModel/TFModelConvert module needs to receive a list of arbitrary Resblock modules.</p>
<p>In the TensorFlow notebook, you will see that it took 58 seconds to train during 5 epochs.<br>
In the PyTorch notebook, it took 122 seconds to train using the same layers and the same number of epochs.</p>
<p><strong>My question is what can I do or what am I doing wrong for PyTorch to be slower than TensorFlow?</strong></p>
<p>TensorFlow notebook: <a href="https://colab.research.google.com/drive/1mG16JPGEPlyZpvYoTOEu4gXPQzE_B47m?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a><br>
PyTorch notebook:  <a href="https://colab.research.google.com/drive/17rqyPeMt0TElw_CSPQ77TuusVHesqDLR?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a></p>

ptrblck: <p>You are recreating randomly initialized layers in each forward pass in <code>ResBlock</code>:</p>
<pre><code class="lang-python">class ResBlock(nn.Module):
    def __init__(self):
        super(ResBlock, self).__init__()
        self.layers = None

    def forward(self, x):
        net = x.clone()

        net = self.layers(net)

        with torch.no_grad():
            # Calculating the kernel size needed to downsample
            out_channels = get_out_channels(self.layers)
            k_size = (x.shape[-2] - net.shape[-2] + 1, x.shape[-1] - net.shape[-1] + 1)

        # Apply an convolution
        conv = nn.Conv2d(x.shape[1], out_channels, k_size).to(DEVICE)

        return net + conv(x)
</code></pre>
<p>which sounds wrong unless you don’t want to train <code>conv</code>.<br>
Initialize the layers in the <code>nn.Module.__init__</code> method and use them in the <code>forward</code> method.</p>

ttiago_jm: <p><code>conv</code> doesn’t need to be trained, it previously had a dirac_ initialization because I just want to reduce the input and don’t want to increase the number of parameters. I could use a MaxPooling but then I would need to create a zero pad tensor, which will be expensive.</p>
<p>But even if I want to train <code>conv</code> how am I going to initialize it in <code>__init__</code> if I need to calculate <code>k_size</code> based on the input form? Is there a way to do this in PyTorch?</p>
<p>To test, I just removed everything from the forward function and wrote this:</p>
<pre><code class="lang-auto">def forward(self, x):
        return self.layers(x)
</code></pre>
<p>The elapsed time didn’t get much better.</p>

ptrblck: <p>Since the conv is not trained at all, you could use adaptive pooling layers instead as I don’t understand how your current approach of using new random filters could work.<br>
In any case, since you still see a slowdown you would need to profile the code using the native profiler or e.g. Nsight Systems and check where the bottlenecks are.<br>
Generally I would also recommend reading through the <a href="https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html">Performance Guide</a>.</p>

ttiago_jm: <p>I have already changed the <code>conv</code> to a pooling layer since the goal is to downsample only.<br>
Thanks for the article! I change a few things and the timing improves a bit.</p>
<p>In a PyTorch group, someone reminded me that PyTorch uses Dynamic Computation Graphs and Tensorflow still uses Static Computation Graphs in the background and is able to do optimizations during training. This is probably why TensorFlow is slightly faster after the first epoch.</p>