有没有办法加快 tf.keras 中的嵌入层？

Question

I'm trying to implement an LSTM model for DNA sequence classification, but at the moment it is unusable because of how long it takes to train (25 seconds per epoch over 6.5K sequences, about 4ms per sample, and we need to train several versions of the model over 100s of thousands of sequences).

DNA sequence can be represented as a string of A, C, G, and T, eg "ACGGGTGACAT" could be an example of a single DNA sequence. Each sequence belongs to one of two categories that I am trying to predict and each sequence contains 1000 characters.

Initially, my model did not include an Embedding layer and instead I manually converted each sequence into a one-hot encoded matrix (4 rows by 1000 columns) and the model didn't work great but was incredibly fast. At this point though I had seen online that using an embedding layer has clear advantages. So I added an embedding layer and instead of using the one-hot encoded matrix I converted the sequences into integers with each character represented by a different integer.

Indeed the model works much better now, but it is about 30 times slower and impossible to work with. Is there something I can do here to speed up the embedding layer?

Here are the functions for constructing and fitting my model:

from tensorflow.keras.layers import Embedding, Dense, LSTM, Activation
from tensorflow.keras import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical

def build_model():
    # initialize a sequential model
    model = Sequential()

    # add embedding layer
    model.add(Embedding(5, 1, input_length=1000, mask_zero=True))

    # Add LSTM layer
    model.add(
       LSTM(5)
    )

    # Add Dense NN layer
    model.add(
        Dense(units=2)
    )

    model.add(Activation('softmax'))

    optimizer = Adam(clipnorm=1.)

    model.compile(
        loss="categorical_crossentropy", optimizer=optimizer, metrics=['accuracy']
    )

    return model

def train_model(X_train, y_train, epochs, batch_size):
    model = build_model()

    # y_train is initially a list of zeroes and ones, needs to be converted to categorical
    y_train = to_categorical(y_train)  

    history = model.fit(
        X_train, y_train, epochs=epochs, batch_size=batch_size
    )

    return model, history

Any help will be greatly appreciated - after much googling and trial-and-error, I can't seem to speed this up.

Answer 1

A possible suggestion is to use a "cheaper" RNN, such as the SimpleRNN instead of LSTM. It has less parameters to train. In some simple testing, I got a ~3x speed up over LSTM, with the same Embedding processing as you currently have. Not sure if you can reduce the sequence length from 1000 to a lower number, but that might be a direction to explore as well. I hope this helps.