Implementation:Tensorflow Tfjs Training Dataset Test

Knowledge Sources	Tensorflow_Tfjs
Domains	Testing, Layers_API
Last Updated	2026-02-10 06:00 GMT

Overview

This is one of the largest test suites in the Layers API (3414 lines), validating LayersModel.fitDataset() and LayersModel.evaluateDataset() methods. These methods enable training and evaluation using TensorFlow.js Dataset objects rather than raw tensors, supporting streaming data pipelines. The tests cover single and multi-input/output models, metrics, validation datasets, callbacks, memory leak verification, batches-per-epoch configuration, class weights with datasets, and various dataset configurations. Reference Python Keras code is included inline for numerical verification.

Code Reference

Source Location: tfjs-layers/src/engine/training_dataset_test.ts (3414 lines)

Repository: GitHub

Test Describe Blocks

• LayersModel.fitDataset - Comprehensive dataset-based training tests:
  • 1 input / 1 output, no metrics, no validation
  • With metrics (accuracy)
  • With validation dataset and validation batches
  • Multi-input and multi-output models
  • Custom callbacks (onBatchEnd, onEpochEnd)
  • Memory leak checks (tfc.memory().numTensors before/after)
  • batchesPerEpoch configuration
  • Early stopping via callbacks
• LayersModel.evaluateDataset - Dataset-based model evaluation

I/O Contract

Inputs to tests:

• FakeNumericDataset instances with configurable shapes, batch sizes, and number of batches
• Dense models compiled with various loss functions and optimizers
• Training configuration: epochs, batchesPerEpoch, validationData, callbacks

Expected outputs/assertions:

• Training history contains expected metric keys and epoch counts
• Loss values match Python Keras reference values (e.g., 0.923649, 0.722993)
• Model weights converge to expected values after training
• No memory leaks: numTensors before training equals numTensors after
• Callbacks receive correct batch/epoch information

Usage Example

describeMathCPUAndGPU('LayersModel.fitDataset', () => {
  it('1 input, 1 output, no metric, no validation, with batchesPerEpoch',
     async () => {
       const model = createDenseModel();
       model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});

       const batchSize = 8;
       const epochs = 2;
       const batchesPerEpoch = 3;
       const dataset = new FakeNumericDataset({
         xShape: [1], yShape: [1], batchSize,
         numBatches: batchesPerEpoch * epochs,
       });

       const history =
           await model.fitDataset(dataset, {batchesPerEpoch, epochs});
       expect(Object.keys(history.history)).toEqual(['loss']);
       expect(history.history.loss.length).toEqual(2);
       expect(history.history.loss[0]).toBeCloseTo(0.923649);
     });
});

Test Coverage Summary

Category	Count	Details
fitDataset Basic	10+	Single I/O, metrics, validation
fitDataset Multi-I/O	5+	Multiple inputs and outputs
fitDataset Callbacks	5+	Custom callbacks, early stopping
fitDataset Memory	5+	Memory leak verification
evaluateDataset	5+	Dataset-based evaluation
Test Environment	CPU and GPU	describeMathCPUAndGPU

Related Pages

• Implementation:Tensorflow_Tfjs_Training_Test
• Implementation:Tensorflow_Tfjs_Training_Weighted_Test
• Implementation:Tensorflow_Tfjs_Base_Callbacks_Test
• Implementation:Tensorflow_Tfjs_Models_Test