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

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
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)
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()
metrics=['accuracy'])
parallel_model = multi_gpu_model(model, gpus=2)

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)


• I would suggest using as little data as possible to help speed up the process. You are right you don't need the database. Although it does help with the problem of organizing the data it may create some latency. – Michael Hearn Nov 15 at 0:42
• Thanks @MichaelHearn. It's true, I should do some learning curves to see how much I need. – J. Montgomery Nov 15 at 3:01

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.