What is the fastest way to train a CNN with billions of examples?

Question

I have a CNN model that I need to train for a large scale genomics application. It is working well with a subset of my training data. I have scaled up to a subset of about 130 million examples and training time is very long, about 3 hours per epoch. I plan to scale up to the hundreds of billions of training examples and I anticipate training time to be not be feasible with my current design. I would appreciate feedback on how I can streamline the training or improve some aspect of my design that I may not be considering. Currently, I am training from a MongoDB. The training examples are not very large. Here is an example.

{
    'added': datetime.datetime(2019, 11, 1, 6, 13, 13, 340000),
    '_id': ObjectId('5dbbccf92464af872756022e'),
    'label': 0,
    'accession': 'GM_0001',
    'data': '34363,30450,9019,19152,8726,22128,59881,17670,15803,64454,64579,28103,52442,64951,29783,64574,652,19243,33498,14775,18803,4700,55446,53912,47645,41465,48257,16305,62071,12334,44698,24371,46515,8445,3000,61849,43228,18120,23587,11105,5453,42707,42739,46122,31285,40773,48162,16653,58783,2928,2836,21330,46947,6719,26992,8852,14520,46212,47362,43554,2147,39372,33885,59716,37384,14825,53387,58763,18065,34070,23278,15641,40237,47950,58811,40015,36880,29841,45351,14904,49660,48224,54638,50358,17202,10701,3564,4829,62655,5684,37207,49724,16369,6769,37827,38144,63885,5070,42882,48960,16178,35758,50554,54253,34556,2383,39431,30176,11482,24459,4472,53825,7764,44500,4869,50875,33037,56353,46848,30769,18729,46026,41409,2826,12092,17086',
    'name': 'Example_1'
}

The relevant data is the 'data' field which is a string of 126 integers where each integer is a value between 0 and about 65,000. The other fields are convenient, but not necessary except for the 'label' field. But even this I could insert into the front of the data field. I mention this because I don't think I necessarily need to train from a MongoDB database.

I am using Keras 2.3.0 with TensorFlow 2.0.0. Below is an example of my code. The workflow is 1) Load a text file containing the document ids of all training examples in the MongoDB collection. I do this so I can shuffle the examples before sending them to the model for training. 2) I load the examples in batches of 50,000 using my Custom_Generator class. This class pulls the documents from the MongoDB using the list of document ids. 3) The model is trained. I use 5 epochs. I currently have 5-fold cross-validation but I know this is not feasible on the full training set. For that I will do a single train-test split. I am currently performing this on a Google Cloud instance with 2 Tesla T4 GPUs. The database is on a bucket. With the cloud I have flexibility of hardware architectures. I would appreciate any insight. This is a rather large engineering challenge for me.

Additional background to the problem: The objective is to classify organisms into broad classes quickly for downstream analysis. The pool of organisms I want to classify is very large (10s of thousands) and very diverse. I'm essentially reading the genomes of the organisms like a book. The genome (a series of "A", "T", "C", or "G") is processed in chunks through a hash function producing integer strings as shown above. Depending on the size of the organism genome, thousands to millions of these integer strings may be produced. So I have many thousands of organisms producing many thousands to millions of examples. To be successful, I feel like I need to capture the diversity of the genomes in the organism pool. To give an example, even though Ecoli and Salmonella are both bacteria, their genomes are quite distinct. I feel like I need to have them both represented in the training set to distinguish them from other organisms I would label as a different class. As far as reducing the dataset, I think I can get by with only training on a representative organism for a give species (since there are many unique genomes available for Ecoli, for example). This will help considerably, but I think the training data set will likely still be in the billions of examples.

import sys
import time
from keras.utils import Sequence, to_categorical, multi_gpu_model
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers import Embedding
from keras.layers.convolutional import Conv1D
from keras.layers.convolutional import MaxPooling1D
from sklearn.model_selection import KFold
from keras.preprocessing.sequence import pad_sequences
import numpy as np
import random
from pymongo import MongoClient
from bson import ObjectId
from sklearn.metrics import classification_report, confusion_matrix


class Custom_Generator(Sequence) :

    def __init__(self, document_ids, batch_size) :
        self.document_ids = document_ids
        self.batch_size = batch_size


    def __len__(self) :
        return (np.ceil(len(self.document_ids) / float(self.batch_size))).astype(np.int)


    def __getitem__(self, idx) :
        client = MongoClient(port=27017)
        db = client[database]
        document_ids = self.document_ids[idx * self.batch_size : (idx+1) * self.batch_size]
        query_results = db[collection].find({'_id': {'$in': document_ids}})
        batch_x, batch_y = [], []
        for result in query_results:
            kmer_list = result['kmers'].split(',')
            label = result['label']
            x = [x for x in kmer_list if len(x) > 0]
            if len(x) < 1:
                continue
            batch_x.append(x)
            one_hot_y = to_categorical(label, 5)
            batch_y.append(one_hot_y)
        batch_x = pad_sequences(batch_x, maxlen=126, padding='post')
        client.close()
        return np.array(batch_x), np.array(batch_y)


# MongoDB database, collection, and document ids of collection
database = 'db'
collection = 'collection_subset2'
docids_file = 'docids_collection_subset2.txt'
id_ls = []
# Convert docids strings to MongoDB ObjectID
with open(docids_file) as f:
    for line in f:
        id_ls.append(ObjectId(line.strip()))
random.shuffle(id_ls)

# Model
model = Sequential()
model.add(Embedding(65521, 100, input_length=126))
model.add(Conv1D(filters=25, kernel_size=5, activation='relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Conv1D(filters=30, kernel_size=3, activation='relu'))
model.add(MaxPooling1D(pool_size=2))
model.add(Flatten())
model.add(Dense(1000, activation='relu'))
model.add(Dense(5, kernel_initializer="normal", activation="softmax"))
metrics=['accuracy'])
parallel_model = multi_gpu_model(model, gpus=2)
parallel_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

seed = 7
batch_size = 50000

# Currently training with 5-fold CV. Will only use single test train split 
# on the full-scale dataset.
kfold = KFold(n_splits=5, shuffle=True, random_state=seed)
kfold_stats = {}
accuracy_ls = []
val_accuracy_ls = []
confusion_ls = []
for fold_idx, (train_idx, test_idx) in enumerate(kfold.split(id_ls)):
    ids_train = np.array(id_ls)[train_idx].tolist()
    ids_test = np.array(id_ls)[test_idx].tolist()
    training_batch_generator = Custom_Generator(ids_train, batch_size)
    validation_batch_generator = Custom_Generator(ids_test, batch_size)
    print('Number of train files: %d' % len(ids_train))
    print('Number of test files: %d' % len(ids_test))
    start = time.time()
    history = parallel_model.fit_generator(
        generator=training_batch_generator,
        steps_per_epoch = int(len(ids_train) // batch_size),
        epochs = 5,
        verbose = 2,
        validation_data = validation_batch_generator,
        validation_steps = int(len(ids_test) // batch_size),
        use_multiprocessing=True
    )
    sys.stderr.write("time to train model (seconds): %d\n"%(time.time() - start))
    sys.stderr.flush()
    print(history.history)
    fold_name = 'kfold_%s' % str(fold_idx)
    kfold_stats.update({fold_name: history.history})
    accuracy_ls.extend(history.history['accuracy'])
    val_accuracy_ls.extend(history.history['val_accuracy'])
    parallel_model.save('model_output_kfold_%s.h5' % str(fold_idx))
    print("Kfold %s finished" % str(fold_idx))
    Y_pred = parallel_model.predict_generator(validation_batch_generator)
    y_pred = np.argmax(Y_pred, axis=1)
    y_true = np.concatenate([np.argmax(batch[1], axis=1) for batch in validation_batch_generator])  
    print('Confusion Matrix')
    conf = confusion_matrix(y_true, y_pred)
    print(conf)
    confusion_ls.append(conf)
    print('Classification Report')
    target_names = ['Class_name_1', 'Class_name_2', 'Class_name_3', 'Class_name_4', 'Class_name_5']
    report = classification_report(y_true, y_pred, target_names=target_names)

Answer 1

First of all, you should add the argument workers = n in the fit generator call. n should be bigger than 1 to prefetch data. As your data processing requires the data be taken from a server or port, you should do pre fetching data as that would fetch the next data while GPU is processing.

If you call fit_generator with workers > 1 , use_multiprocessing=True , we will prefetch queue_size batches.

Source: https://github.com/keras-team/keras/issues/12847

Secondly, as @MichaelHearn mentioned, you should plot a graph to see how many data is needed. You task should not require several billion data samples. A couple hundred thousand data should be enough. If the model is not getting a good accuracy from a couple hundred thousand data samples, perhaps you can try changing up the model architecture instead of adding more data. Maybe you can add drop out layers.

Hope I can help you.