How to speed up recurrent autoencoder?

Cola_Lightyear: <p>I ported my previous code to pytorch and now am looking for possibilities to speed up a recurrent autoencoder, which greatly impacts the training (and to a lesser degree the inference) speed.<br>
This is the forward call (+ additional functions) of my recurrent autoencoder. The for loop in the forward call slows down the processing</p>
<p>Recurrent autoencoder code:</p>
<pre><code class="lang-auto">class FRAE(nn.Module):
    def __init__(self, recurrent_timesteps=1, neurons_per_layer=[22, 30, 14, 5], hidden_size=22, quant=None):
        super().__init__()

        self.hidden_size = hidden_size * recurrent_timesteps
        self.recurrent_timesteps = recurrent_timesteps
        self.neurons_per_layer = neurons_per_layer
        self.quant = quant

        # build layers according to list neurons_per_layer
        # encoder
        layers_enc = []
        for i in range(len(neurons_per_layer) - 1):

            # first layer takes input and hidden vector
            if (i == 0):
                layers_enc.append(nn.Linear(neurons_per_layer[i] + self.hidden_size, neurons_per_layer[i + 1]))
            else:
                layers_enc.append(nn.Linear(neurons_per_layer[i], neurons_per_layer[i + 1]))
            layers_enc.append(nn.SiLU())

        # decoder
        neurons_per_layer_rev = neurons_per_layer[::-1]
        layers_dec = []
        for i in range(len(neurons_per_layer) - 1):
            # first layer takes input and hidden vector
            if (i == 0):
                layers_dec.append(nn.Linear(neurons_per_layer_rev[i] + self.hidden_size, neurons_per_layer_rev[i + 1]))
                layers_dec.append(nn.SiLU())
            else:
                layers_dec.append(nn.Linear(neurons_per_layer_rev[i], neurons_per_layer_rev[i + 1]))
                if (i == len(neurons_per_layer) - 2):
                    pass # linear activation
                    # layers_dec.append(nn.Tanh())
                else:
                    layers_dec.append(nn.SiLU())

        self.encoder = nn.Sequential(*layers_enc)
        self.decoder = nn.Sequential(*layers_dec)


    def processFrame(self, x, hidden):
        combined_enc = torch.cat((x, hidden), dim=1)
        encoded = self.encoder(combined_enc)

        if self.quant:
            encoded = encoded.detach().numpy()
            # encoded = np.reshape(encoded, (1, self.neurons_per_layer[-1]))
            encoded = self.quant.encode(encoded)
            encoded = self.quant.decode(encoded)
            encoded = torch.from_numpy(encoded)
            # encoded = torch.reshape(encoded, (-1,))

        combined_dec = torch.cat((encoded, hidden), dim=1)
        decoded = self.decoder(combined_dec)

        new_hidden = self.updateHidden(decoded, hidden)
        return decoded, new_hidden


    def initHidden(self, inputs):
        return torch.zeros((inputs.shape[0], self.hidden_size)).to(inputs.device)

    def updateHidden(self, last_decoded, hidden):
        if self.recurrent_timesteps == 1:
            new_hidden = last_decoded
        else:
            hid_size = self.hidden_size // self.recurrent_timesteps
            new_hidden = torch.zeros_like(hidden)
            new_hidden[:, hid_size:] = hidden[:, :self.hidden_size - hid_size]
            new_hidden[:, :hid_size] = last_decoded

        return new_hidden

    def forward(self, inputs): # 64 x 64 x 7999 or 64 x 512 x 7999
        if len(inputs.shape) == 2:  # no batch
            inputs = torch.unsqueeze(inputs, 0)

        recon = torch.zeros(inputs.size()).to(inputs.device)
        new_hidden = self.initHidden(inputs)
        for i in range(inputs.size()[-1]): # this is slow on cpu
            recon_i, new_hidden = self.processFrame(inputs[:, :, i], new_hidden)
            recon[:, :, i] = recon_i

        return recon

</code></pre>
<p>The input dimension of the forward method is currently 64 x 64 x 7999 and 64 x 512 x 7999. Unfortunately, I cannot reduce the amount of data and the input dimensionality I use.</p>
<p>I am using four recurrent autoencoders in total, each with about 70.000 to 1.2 million parameters.The recurrent autoencoders are part of a larger model with in total about 3.5 million parameters. The recurrent autoencoders make up 3 million of the 3.5 million parameters.</p>
<p>The rest of the larger model does consists of stacks of convolutional layers, so is not recurrent. Without the recurrent autoencoders, the model trainings quickly. I tried moving the autoencoders to cpu, which can help with the for loop, but unfortunately this turned out considerably slower than before.</p>
<p>Is there some way to make this train faster? Currently a batch requires about 85 seconds to finish, which is way too slow for development. The forward pas requires about 13 seconds, however,  loss.backward() requires about 70 seconds, so is really slow.</p>
<p>Because the amount of data is so large, I cannot increase the batch size above about 16 without running out of memory.</p>