Merge pull request #20 from RonasIT/14-asynchronously-loaded-datasets

Asynchronously loaded datasets

Author: astorozhevsky

README.md

@@ -127,6 +127,51 @@ await trainer.trainAndTest({
 });
 ```
 
+##### Loading data asynchronously
+
+When working with large dataset, you might find out that the whole dataset
+can't fit in memory. In this situation you might want to load the data in
+chunks. To do this, you can define the asynchronous generators for
+`trainingDataset`, `validationDataset` and `testingDataset`.
+
+This library provides the `makeChunkedDataset` helper to make it easier to
+create chunked datasets where chunks are controlled with `skip` and `take`
+parameters.
+
+`makeChunkedDataset` helper accepts the following parameters:
+
+- `loadChunk` – an asynchronous function accepting the numeric `skip` and `take`
+  parameters and returning an array of samples.
+- `chunkSize` – the number of samples loaded per chunk.
+- `batchSize` – the number of samples in each batch.
+
+```typescript
+const loadTrainingSamplesChunk = async (skip: number, take: number): Promise<Array<Sample>> => {
+  // Your samples chunk loading logic goes here. For example, you may want to
+  //   load samples from database, or from a remote data source.
+};
+
+const makeTrainingDataset = (): data.Dataset<TensorContainer> => makeChunkedDataset({
+  loadChunk: loadTrainingSamplesChunk,
+  chunkSize: 32,
+  batchSize: 32
+});
+
+// You should also define similar functions for validationDataset and
+//   trainingDataset. We omit this for the sake of brevity.
+
+const trainingDataset = makeTrainingDataset();
+const validationDataset = makeValidationDataset();
+const testingDataset = makeTestingDataset();
+
+await trainer.trainAndTest({
+  trainingDataset,
+  validationDataset,
+  testingDataset,
+  printTestingResults: true
+});
+```
+
 #### Saving the model
 
 To save the trained model, you need to call the `save` method of the

packages/tfjs-node-helpers/src/feature-engineering/extract-features.ts

@@ -0,0 +1,32 @@
+import { Sample } from '../training/sample';
+import { FeatureExtractor } from './feature-extractor';
+
+export const extractFeatures = async <D, T>({
+  data,
+  inputFeatureExtractors,
+  outputFeatureExtractor
+}: {
+  data: Array<D>;
+  inputFeatureExtractors: Array<FeatureExtractor<D, T>>;
+  outputFeatureExtractor: FeatureExtractor<D, T>;
+}): Promise<Array<Sample>> => {
+  const samples = [];
+
+  for (const dataItem of data) {
+    const [inputFeatures, outputFeature] = await Promise.all([
+      Promise.all(
+        inputFeatureExtractors.map((featureExtractor) => {
+          return featureExtractor.extract(dataItem);
+        })
+      ),
+      outputFeatureExtractor.extract(dataItem)
+    ]);
+
+    const input = inputFeatures.map((feature) => feature.value);
+    const output = [outputFeature.value];
+
+    samples.push({ input, output });
+  }
+
+  return samples;
+}

packages/tfjs-node-helpers/src/feature-engineering/prepare-datasets-for-binary-classification.ts

@@ -1,6 +1,7 @@
 import { data, TensorContainer } from '@tensorflow/tfjs-node';
 import { splitSamplesIntoTrainingValidationTestForBinaryClassification } from '../data-splitting/training-validation-test-for-binary-classification';
-import { makeDataset } from '../utils/make-dataset-from-array';
+import { makeDataset } from '../utils/make-dataset';
+import { extractFeatures } from './extract-features';
 import { FeatureExtractor } from './feature-extractor';
 
 export const prepareDatasetsForBinaryClassification = async <D, T>({
@@ -24,23 +25,11 @@ export const prepareDatasetsForBinaryClassification = async <D, T>({
   validationDataset: data.Dataset<TensorContainer>;
   testingDataset: data.Dataset<TensorContainer>;
 }> => {
-  const samples = [];
-
-  for (const dataItem of data) {
-    const [inputFeatures, outputFeature] = await Promise.all([
-      Promise.all(
-        inputFeatureExtractors.map((featureExtractor) => {
-          return featureExtractor.extract(dataItem);
-        })
-      ),
-      outputFeatureExtractor.extract(dataItem)
-    ]);
-
-    const input = inputFeatures.map((feature) => feature.value);
-    const output = [outputFeature.value];
-
-    samples.push({ input, output });
-  }
+  const samples = await extractFeatures({
+    data,
+    inputFeatureExtractors,
+    outputFeatureExtractor
+  });
 
   const { trainingSamples, validationSamples, testingSamples } = splitSamplesIntoTrainingValidationTestForBinaryClassification(
     samples,

packages/tfjs-node-helpers/src/index.ts

@@ -1,7 +1,9 @@
 export * from './classification/binary-classification-trainer';
 export * from './classification/binary-classifier';
 export * from './data-splitting/training-validation-test-for-binary-classification';
+export * from './feature-engineering/extract-features';
 export * from './feature-engineering/feature-extractor';
 export * from './feature-engineering/feature';
 export * from './feature-engineering/prepare-datasets-for-binary-classification';
 export * from './training/sample';
+export * from './utils/make-chunked-dataset';

packages/tfjs-node-helpers/src/utils/make-chunked-dataset-generator.ts

@@ -0,0 +1,30 @@
+import { Tensor, tensor } from '@tensorflow/tfjs-node';
+import { Sample } from '../training/sample';
+
+export const makeChunkedDatasetGenerator = async function* ({
+  loadChunk,
+  chunkSize
+}: {
+  loadChunk: (skip: number, take: number) => Promise<Array<Sample>>,
+  chunkSize: number
+}): AsyncGenerator<{ xs: Tensor, ys: Tensor }> {
+  let skip = 0;
+  let take = chunkSize;
+
+  while (true) {
+    const samples = await loadChunk(skip, take);
+
+    for (const sample of samples) {
+      yield {
+        xs: tensor(sample.input),
+        ys: tensor(sample.output)
+      };
+    }
+
+    if (samples.length < take) {
+      break;
+    }
+
+    skip += take;
+  }
+};

packages/tfjs-node-helpers/src/utils/make-chunked-dataset.ts

@@ -0,0 +1,22 @@
+import { data, TensorContainer } from '@tensorflow/tfjs-node';
+import { Sample } from '../training/sample';
+import { makeChunkedDatasetGenerator } from '../utils/make-chunked-dataset-generator';
+
+export const makeChunkedDataset = ({
+  loadChunk,
+  chunkSize,
+  batchSize
+}: {
+  loadChunk: (skip: number, take: number) => Promise<Array<Sample>>,
+  chunkSize: number,
+  batchSize: number
+}): data.Dataset<TensorContainer> => {
+  return data
+    .generator(
+      () => makeChunkedDatasetGenerator({
+        loadChunk,
+        chunkSize
+      }) as any
+    )
+    .batch(batchSize);
+};

…ers/src/utils/make-dataset-from-array.ts …s-node-helpers/src/utils/make-dataset.tspackages/tfjs-node-helpers/src/utils/make-dataset-from-array.ts renamed to packages/tfjs-node-helpers/src/utils/make-dataset.ts packages/tfjs-node-helpers/src/utils/make-dataset-from-array.ts renamed to packages/tfjs-node-helpers/src/utils/make-dataset.ts

@@ -3,7 +3,8 @@
   "compilerOptions": {
     "outDir": "../../dist/out-tsc",
     "declaration": true,
-    "types": []
+    "types": [],
+    "lib": ["es2018"]
   },
   "include": ["**/*.ts"],
   "exclude": ["jest.config.ts", "**/*.spec.ts", "**/*.test.ts"]