Repository: pytorch
Issue: Conv-RNN combination slow in backward pass
Labels: module: performance,module: nn,triaged

Models used for processing audio signals often combine a Conv layer with a GRU/LSTM layer. In pytorch, this kind of combination seems to yield an unexpectedly long backward pass. Here is a minimal model of this kind:

```python
class PyTorchModel(nn.Module):
    def __init__(self, num_channels, kernel_size):
        super().__init__()

        self.conv = nn.Conv1d(num_channels, num_channels, kernel_size, padding=int((kernel_size - 1) / 2))
        self.gru = nn.GRU(num_channels, num_channels, batch_first=True)
        self.fc = nn.Sequential(
            nn.Linear(num_channels, 1),
            nn.Sigmoid()
        )

    def forward(self, inp, do_conv=True, do_gru=True):
        out = self.conv(inp)
        # (batchsize, num_channels, seq_length) -> (batchsize, seq_length, num_channels)
        out = out.transpose(1, 2)
        out, _ = self.gru(out)
        out = self.fc(out)
        return out
```

The script provided at the end benchmarks this model and compares the execution times to keras (with tensorflow backend). Here are the benchmark results (obtained on a MacBook Pro with the CPU i7-4980HQ):

```
data for pytorch: X.shape=(32, 128, 2048), Y.shape=(32, 2048, 1)
data for keras: X.shape=(32, 2048, 128), Y.shape=(32, 2048, 1)

>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
Testing Conv and GRU together
pytorch: forward_pass=1.3s backward_pass=40.0s
keras: forward_pass=0.7s backward_pass=2.5s
<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
Testing Conv only
pytorch: forward_pass=0.5s backward_pass=0.2s
keras: forward_pass=0.3s backward_pass=0.6s
<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
Testing GRU only
pytorch: forward_pass=0.8s backward_pass=0.9s
keras: forward_pass=0.4s backward_pass=1.8s
<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
```

The backward pass for the Conv-GRU model takes 16 times longer in pytorch as compared to keras. There is no such strong difference for a model with only a Conv or a GRU layer. Hence the problems seems to stem from the combination. 

An apparent difference between the pytorch and keras implementations is that in pytorch there is a transpose operation between the Conv and the GRU layers. My initial guess was that this is the source of the problem. However, in the benchmark above the cases "Conv only" and "GRU only" also contain a transpose operation (see benchmark script), which doesn't seem to influence the performance negatively.

The results are similar when the GRU is replaced by a LSTM or RNN layer.

Information about the used setup:
- OS: OSX (but tested also on Linux, with similar results)
- PyTorch version: 0.3.1.post2
- How you installed PyTorch: `conda install pytorch -c pytorch `
- Python version: 3.6.3
- CUDA/cuDNN version: no CUDA installed

Script used for benchmarking:
```python
import time

import torch
from torch.autograd import Variable
from torch import nn

import keras.models
from keras.layers import Input, Conv1D, GRU, Dense, TimeDistributed
from keras.optimizers import SGD


batchsize = 32
seq_length = 2048
num_channels = 128
kernel_size = 15


def generate_training_data():
    # for pytorch
    X = torch.randn((batchsize, num_channels, seq_length))
    Y = torch.rand((batchsize, seq_length, 1)).round()

    # for keray
    Xnp = X.transpose(1, 2).clone().numpy()
    Ynp = Y.clone().numpy()

    return X, Y, Xnp, Ynp


class PyTorchModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv1d(num_channels, num_channels, kernel_size, padding=int((kernel_size - 1) / 2))
        self.gru = nn.GRU(num_channels, num_channels, batch_first=True)
        self.fc = nn.Sequential(
            nn.Linear(num_channels, 1),
            nn.Sigmoid()
        )

    def forward(self, inp, do_conv=True, do_gru=True):
        if do_conv:
            out = self.conv(inp)
        else:
            out = inp

        # (batchsize, num_channels, seq_length) -> (batchsize, seq_length, num_channels)
        out = out.transpose(1, 2)

        if do_gru:
            out, _ = self.gru(out)

        out = self.fc(out)
        return out


def get_keras_model(do_conv, do_gru):
    inp = Input((seq_length, num_channels))
    out = inp

    if do_conv:
        out = Conv1D(num_channels, kernel_size, padding="same")(out)

    if do_gru:
        out = GRU(num_channels, return_sequences=True)(out)

    out = TimeDistributed(Dense(1, activation="sigmoid"))(out)

    return keras.models.Model(inputs=inp, outputs=out)


def benchmark_pytorch(X, Y, do_conv, do_gru, num_iter=5):
    """
    Do the foward and backward pass `num_iter` times and return the average runtime.
    """
    loss_fn = nn.BCELoss()
    model = PyTorchModel()

    forward_time = 0
    backward_time = 0

    for _ in range(num_iter):
        # do forward pass
        start = time.time()
        out = model(Variable(X), do_gru=do_gru, do_conv=do_conv)
        forward_time += time.time() - start

        loss = loss_fn(out, Variable(Y))

        # do backward pass
        start = time.time()
        loss.backward()
        backward_time += time.time() - start

    return forward_time / num_iter, backward_time / num_iter


def benchmark_keras(X, Y, do_conv, do_gru, num_iter=5):
    """
    Do the foward and backward pass `num_iter` times and return the average runtime.
    """
    forward_time = 0
    backward_time = 0

    model = get_keras_model(do_conv=do_conv, do_gru=do_gru)
    model.compile(loss='binary_crossentropy', optimizer=SGD(lr=0.001))

    for _ in range(num_iter):
        # do forward pass
        start = time.time()
        model.predict(X)
        forward_time += time.time() - start

        # do backward pass
        start = time.time()
        model.train_on_batch(X, Y)
        backward_time += time.time() - start

    return forward_time / num_iter, backward_time / num_iter


if __name__ == '__main__':
    X, Y, Xnp, Ynp = generate_training_data()

    print("")
    print("data for pytorch: X.shape={}, Y.shape={}".format(tuple(X.shape), tuple(Y.shape)))
    print("data for keras: X.shape={}, Y.shape={}".format(tuple(Xnp.shape), tuple(Ynp.shape)))

    parameters = [
        (True, True, "Testing Conv and GRU together"),
        (True, False, "Testing Conv only"),
        (False, True, "Testing GRU only"),
    ]

    for do_conv, do_gru, description in parameters:
        print("\n" + ">" * 50)
        print(description)

        print("pytorch: forward_pass={:.1f}s backward_pass={:.1f}s".format(
            *benchmark_pytorch(X, Y, do_conv=do_conv, do_gru=do_gru)
        ))
        print("keras: forward_pass={:.1f}s backward_pass={:.1f}s".format(
            *benchmark_keras(Xnp, Ynp, do_conv=do_conv, do_gru=do_gru)
        ))

        print("<" * 50)
```