Training on Large Datasets

Neural nets are well-suited for being trained on very large datasets, even those that are too large to fit into memory. The most popular optimization algorithms for training neural nets (such as "ADAM" or "RMSProp" in (RefLink[NetTrain, paclet:ref/NetTrain])) are variations of an approach called stochastic gradient descent. In this approach, small batches of data are randomly sampled from the full training dataset and used to perform a parameter update. Thus, neural nets are an example of an online learning algorithm, which does not require the entire training dataset to be in memory. This is in contrast to methods such as the Support Vector Machine (SVM) and Random Forest algorithms, which usually require the entire dataset to reside in memory during training.

However, special handling is required if (RefLink[NetTrain, paclet:ref/NetTrain]) is to be used on a dataset that does not fit into memory, as the full training dataset cannot be loaded into a Wolfram Language session. There are two approaches to training on such large datasets. The first approach is for users to write a generator function f that, when evaluated, can load a single batch of data from an external source such as a disk or database. (RefLink[NetTrain, paclet:ref/NetTrain])[net,f,…] calls f at each training batch iteration, thus only keeping a single batch of training data in memory. This is the most flexible approach to doing out-of-core learning. However, it is up to the user to ensure that the evaluation of f is fast enough to avoid slowing down net training (which is particularly important when training on extremely powerful GPUs) and that it is correctly sampling from the full training dataset.

The second approach works in the special case of a training dataset consisting of image or audio files. In this case, encoders such as the (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) can efficiently read and preprocess images stored on disk.

Find a file in (RefLink[ExampleData, paclet:ref/ExampleData]):

f=File[FindFile["ExampleData/coneflower.jpg"]]

Apply the encoder to the file:

NetEncoder[{"Image",{2,2}}][f]

(RefLink[NetTrain, paclet:ref/NetTrain])[net,{(RefLink[File, paclet:ref/File])[…]->…,…},…] then automatically does out-of-core learning on a dataset where (RefLink[File, paclet:ref/File])[…] represents an out-of-core image or audio file.

Some general performance tips for doing out-of-core learning:

Where possible, use (RefLink[NetEncoder, paclet:ref/NetEncoder]), whose implementation has been highly optimized.

(RefLink[NetTrain, paclet:ref/NetTrain]) loads the next training data batch in parallel with doing a training iteration on the previous training data batch. Thus, the data loading and preprocessing only have an impact on training speed if it is slower than doing a single training iteration. This is highly dependent on the complexity of the net and the hardware the net is being trained on.

When using a generator function f, ensure it is fast enough to avoid slowing down training. To do this, use the generator function to create a copy of the training data that will fit into memory using e.g. dataset=f[<|"BatchSize"->1000|>], and measure the speed at which (RefLink[NetTrain, paclet:ref/NetTrain]) operates on this data using (RefLink[NetTrain, paclet:ref/NetTrain])[net,dataset,"MeanInputsPerSecond",(RefLink[TimeGoal, paclet:ref/TimeGoal])->30]. Then, compare this to the speed when using the data generator using (RefLink[NetTrain, paclet:ref/NetTrain])[net,f,"MeanInputsPerSecond",TimeGoal->30]. If the generator is slower, then either optimize f or accept this slowdown as the price of doing out-of-core training.

Solid State Drives (SSDs) typically give much better performance than Hard Disk Drives (HDDs) when reading a random batch of files from disk. This is due to the random seek time being around an order of magnitude faster on SSDs versus HDDs.

If the audio sampling rate of audio or the image dimensions of image files are significantly different from what the net is being trained on, significant training slowdowns can occur, as the expensive resampling operation is not done only once. The best approach is to create a new set of resampled image or audio files and train on these.

Use an NVIDIA GPU for training. This is done using the (RefLink[TargetDevice, paclet:ref/TargetDevice]) option in (RefLink[NetTrain, paclet:ref/NetTrain]). Training on a GPU often speeds up training by an order of magnitude compared to only using a CPU. The reason for this is that net evaluations are typically much faster on GPUs. In addition, data loading and preprocessing can be faster, as the CPU can focus completely on this task.

Out-of-Core Training on MNIST

This example shows how to train a net on the MNIST dataset when the images are stored on disk as (StringTypeLink["JPEG", paclet:ref/format/JPEG]) files rather than stored in memory. Although it is hardly necessary to training MNIST, this method can be used to train nets on terabyte-scale image datasets such as the ImageNet dataset.

We are able to use the special syntax (RefLink[NetTrain, paclet:ref/NetTrain])[net,{(RefLink[File, paclet:ref/File])[…]->(Subscript[class, 1]),…},…] as long as we attach an (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) to the input port of net. We also show how to reproduce this using the more flexible (RefLink[NetTrain, paclet:ref/NetTrain]) generator function syntax.

Preparing the Data

We first need to create an out-of-core version of the MNIST that is of the form {(RefLink[File, paclet:ref/File])[…]->(Subscript[class, 1]),…}.

Obtain the MNIST training and test sets:

trainingData=ResourceData["MNIST","TrainingData"];
testData=ResourceData["MNIST","TestData"];

RandomSample[trainingData,5]

Let us save the images as (StringTypeLink["JPEG", paclet:ref/format/JPEG]) files in the default temporary directory (RefLink[$TemporaryDirectory, paclet:ref/$TemporaryDirectory]). Each exported file also requires a unique name. A good approach to creating unique names is to use the (RefLink[Hash, paclet:ref/Hash]) function, which returns a unique hash for every unique image.

Define a function that exports an image to a (StringTypeLink["JPEG", paclet:ref/format/JPEG]) file in the (RefLink[$TemporaryDirectory, paclet:ref/$TemporaryDirectory]):

exportImage[img_]:=File@Export[FileNameJoin[{$TemporaryDirectory,Hash[img,Automatic,"HexString"]<>".jpg"}],img]

Export an image:

exportImage[]

When training on images, it is usually the case that the net requires the images to be a single size, color space, etc. In the case of LeNet, the images are expected to be grayscale and of dimensions 28×28. It will usually increase the speed of training if the images that will be read from disk are already conformed to what the net expects, rather than needing to conform the images every time they are loaded from disk. If images are not already conformed, modifying exportImage to conform the images using (RefLink[ConformImages, paclet:ref/ConformImages]) is recommended.

We can now map the exporting function across the training and test sets. One optimization is to do the exporting in parallel, using (RefLink[ParallelMap, paclet:ref/ParallelMap]) rather than (RefLink[Map, paclet:ref/Map]).

Export the images and create new training and test sets of the form {(RefLink[File, paclet:ref/File])[…]->class,…}:

trainingDataFiles=ParallelMap[exportImage[First@#]->Last[#]&,trainingData];

testDataFiles=ParallelMap[exportImage[First@#]->Last[#]&,testData];

Display a (RefLink[RandomSample, paclet:ref/RandomSample]) of the new training set:

RandomSample[trainingDataFiles,3]

Now only the references to the image files need to be kept in memory, which is much smaller than keeping the images themselves in memory.

Obtain the (RefLink[ByteCount, paclet:ref/ByteCount]) of the two training sets:

ByteCount[trainingData]

ByteCount[trainingDataFiles]

Simple Out-of-Core Training

Define a convolutional neural network that has an (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) attached to the input port:

lenet=NetChain[{
ConvolutionLayer[20,5],Ramp,PoolingLayer[2,2],
ConvolutionLayer[50,5],Ramp,PoolingLayer[2,2],
FlattenLayer[],500,Ramp,10,SoftmaxLayer[]},
"Output"->NetDecoder[{"Class",Range[0,9]}],
"Input"->NetEncoder[{"Image",{28,28},"Grayscale"}]
]

As this net has an (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) attached to the input port, it can take images represented as (RefLink[File, paclet:ref/File]) objects as input.

Obtain a file from the training set:

testFile=trainingDataFiles[[1,1]]

Initialize the net with (RefLink[NetInitialize, paclet:ref/NetInitialize]) and apply it to the file:

NetInitialize[lenet][testFile]

Training the net on data of the form {(RefLink[File, paclet:ref/File])[…]->(Subscript[class, 1]),…} is exactly the same as training on data of the form {(Subscript[image, 1])->(Subscript[class, 1]),…}.

Train the net for three training rounds:

results=NetTrain[lenet,trainingDataFiles,All,ValidationSet->testDataFiles,MaxTrainingRounds-> 3]
trained=results["TrainedNet"]

Evaluate the trained network directly on images randomly sampled from the validation set:

imgs=Keys@RandomSample[testData,5];
Thread[imgs->trained[imgs]]

Obtain the accuracy of the trained net on the test set:

NetMeasurements[trained,testDataFiles,"Accuracy"]

Training with a Generator Function

We can also use the more general generator syntax for (RefLink[NetTrain, paclet:ref/NetTrain]). This approach is more flexible, allowing custom image preprocessing, data augmentation, etc.

Define a generator function that returns batches of training data:

genTrain=Function[RandomSample[trainingDataFiles,#BatchSize]]

Generate a batch of 4 training examples:

genTrain[<|"BatchSize"->4|>]

Train the net with the generator function:

results=NetTrain[lenet,{genTrain,"RoundLength"->Length[trainingDataFiles]},All,ValidationSet->testDataFiles,MaxTrainingRounds-> 3]
trained2=results["TrainedNet"]

Obtain the accuracy of the trained net on the test set:

NetMeasurements[trained2,testDataFiles,"Accuracy"]

Import Versus NetEncoder Performance

Using a (RefLink[NetEncoder, paclet:ref/NetEncoder]) for loading data is usually much faster than writing a data loader using top-level Wolfram Language code. As an example, let us compare a simple image importer to the (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]).

Define a custom image data loader:

customImageEnc[files_]:=Map[Import/*ImageData/*List,files]

Define an (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]):

imageEnc=NetEncoder[{"Image",28,"ColorSpace"->"Grayscale"}]

Each of these encoders will produce a rank-4 array when applied to a list of input files.

Create a batch of 4 files:

fileList=Table[trainingDataFiles[[1,1]],4]

Apply both image encoders to a list of files and obtain the output dimensions and time taken to evaluate:

timing1=RepeatedTiming@Dimensions[customImageEnc[fileList]]

timing2=RepeatedTiming@Dimensions[imageEnc[fileList]]

Obtain the speedup of the (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) compared to the custom image encoder:

First[timing1]/First[timing2]

As can be seen, the (StringTypeLink["Image", paclet:ref/netencoder/Image]) (RefLink[NetEncoder, paclet:ref/NetEncoder]) is more than 400 times faster than a custom image encoder! The difference is even greater if the batch size is increased.

Using a MongoDB Database

A (WebLink[MongoDB database, https://www.mongodb.com/]) is one solution for storing large datasets. The Wolfram Language has the MongoLink package for interacting with MongoDB databases.

This example shows how to train a net on the toy Fisher Iris dataset stored in a MongoDB database. Only small batches of data will be randomly sampled from the database during each training iteration. Thus, this method scales to datasets that cannot be stored in memory.

We will assume you are familiar with MongoLink and MongoDB databases. If not, reading the MongoLink Introduction tutorial is recommended before continuing.

It is also assumed that a MongoDB server is running on your local machine at the default host and port. For platform-dependent instructions for running a MongoDB server locally, see (WebLink[this, https://docs.mongodb.com/manual/installation/]).

Data Insertion

The dataset we are using is the Fisher Iris dataset, where the task is to classify a flower into four classes based on some numerical features.

Obtain the Fisher Iris dataset:

trainingData=ExampleData[{"MachineLearning","FisherIris"},"TrainingData"];
testData=ExampleData[{"MachineLearning","FisherIris"},"TestData"];

RandomSample[trainingData,5]

Now we insert the training data into a MongoDB database.

Load MongoLink:

Needs["MongoLink`"]

Create a client connection using the default host "localhost" and port 27017 (this is the default hostname and port when running the MongoDB server on your local machine):

client=MongoConnect[]

Let us create a MongoDB collection named "WolframNNTutorialCollection" in the database "WolframNNTutorialDatabase".

Create a collection using (PackageLink[MongoGetCollection, paclet:MongoLink/ref/MongoGetCollection, MongoLink Package Symbol]):

collection=MongoGetCollection[client,"WolframNNTutorialDatabase","WolframNNTutorialCollection"]

If the collection and database do not yet exist, they will be created when we first insert data into them.

Convert the training data into a list of MongoDB documents:

trainingDataDoc=<|"Input"->First[#],"Output"->Last[#]|>&/@trainingData;

RandomSample[trainingDataDoc,2]

Insert the training data into the collection:

MongoCollectionInsert[collection,trainingDataDoc]

Take a random sample of the collection to verify that the data was inserted:

RandomSample[collection,2]

Build a Classification Net

Now let us build a net that can perform classification on the dataset. We first require the list of all possible classes, which is best generated by a database query.

Build a list of the unique labels to which each example is assigned using (PackageLink[MongoCollectionDistinct, paclet:MongoLink/ref/MongoCollectionDistinct, MongoLink Package Symbol]):

labels=MongoCollectionDistinct[collection,"Output"]

Create a (RefLink[NetChain, paclet:ref/NetChain]) to perform the regression, using a "Class" decoder to interpret the output of the net as probabilities for each class:

net=NetChain[{3,SoftmaxLayer[]},"Input"->4,"Output"->NetDecoder[{"Class",labels}]]

Construct a Generator Function

We need to define a generator function to use in (RefLink[NetTrain, paclet:ref/NetTrain]). This function needs to randomly sample documents from the "WolframNNTutorialCollection".

Use (PackageLink[MongoCollectionAggregate, paclet:MongoLink/ref/MongoCollectionAggregate, MongoLink Package Symbol]) with the (WebLink["$sample", https://docs.mongodb.com/manual/reference/operator/aggregation/sample/#pipe._S_sample]) aggregation operator to obtain two random samples from the "WolframNNTutorialCollection" collection:

ReadList@MongoCollectionAggregate[collection,{<|"$sample"-><|"size"->2|>|>}]

This is equivalent to using (RefLink[RandomSample, paclet:ref/RandomSample]) on the collection:

RandomSample[collection,2]

However, we do not need the "_id" field. We can remove it either by modifying the Wolfram Language output or, more efficiently, by adding the (WebLink["$project", https://docs.mongodb.com/manual/reference/operator/aggregation/project/#pipe._S_project]) operator to the aggregation pipeline.

Read from the collection with the "_id" field removed:

ReadList@MongoCollectionAggregate[collection,{<|"$project"-><|"_id"->False|>|>,<|"$sample"-><|"size"->2|>|>}]

Define a generator function for use in (RefLink[NetTrain, paclet:ref/NetTrain]):

gen1=Function[ReadList@MongoCollectionAggregate[collection,{<|"$project"-><|"_id"->False|>|>,<|"$sample"-><|"size"->#BatchSize|>|>}]];

This generator function correctly returns a randomly sampled list of examples when the batch size is specified:

gen1[<|"BatchSize"->2|>]

There are two main valid forms of training data that generator functions can produce: either a list of example associations {<|(Subscript[key, 1])->(Subscript[val, 11]),(Subscript[key, 2])->(Subscript[val, 21]),…|>,<|(Subscript[key, 1])->(Subscript[val, 12]),…|>,…} produced by the generator, or a single (RefLink[Association, paclet:ref/Association]) where each key has a list of example values <|(Subscript[key, 1])->{(Subscript[val, 11]),(Subscript[val, 12]),…},(Subscript[key, 2])->{(Subscript[val, 21]),(Subscript[val, 22]),…},…|>. One form of the generator function output out can be converted to the other via (RefLink[Transpose, paclet:ref/Transpose])[out,AllowedHeads->(RefLink[All, paclet:ref/All])].

MongoDB can directly produce the second form of grouping examples together using the (WebLink[$group, https://docs.mongodb.com/manual/reference/operator/aggregation/group/#pipe._S_group]) aggregation stage. This is often much more efficient than producing the first form.

Define a generator function that returns batches of examples grouped together:

gen2=Function[First@ReadList@MongoCollectionAggregate[collection,{<|"$sample"-><|"size"->#BatchSize|>|>,<|"$group"-><|"_id"->Null,"Input"-><|"$push"->"$Input"|>,"Output"-><|"$push"->"$Output"|>|>|>,<|"$project"-><|"_id"->False|>|>}]];

Generate a batch of two examples using the second generator:

gen2[<|"BatchSize"->2|>]

The grouped generator is faster at producing batches of size 64:

gen1[<|"BatchSize"->64|>];//RepeatedTiming

gen2[<|"BatchSize"->64|>];//RepeatedTiming

Obtain the samples generated per second by the second generator:

64/First[%]

Train the Net

Train the net with the generator function:

results=NetTrain[net,gen2,All,BatchSize->64,ValidationSet->testData]
trained=results["TrainedNet"]

There is an issue with this approach: the performance on the test set is computed after every batch is trained on, compared to the usual case where it is computed after a single pass through the entire training dataset (a round). When using a generator, (RefLink[NetTrain, paclet:ref/NetTrain]) has no idea what the size of a round should be, unless we specify it explicitly.

Obtain the total number of examples in the collection using (PackageLink[MongoCollectionCount, paclet:MongoLink/ref/MongoCollectionCount, MongoLink Package Symbol]):

roundSize=MongoCollectionCount[collection]

Train the net with a specified round size for 2000 rounds:

results=NetTrain[net,{gen2,"RoundLength"->roundSize},All,BatchSize->64,ValidationSet->testData,MaxTrainingRounds->2000]
trained=results["TrainedNet"]

Test the accuracy of the trained net on the test set:

NetMeasurements[trained,testData,"Accuracy"]