Mutable.ai logoAuto Wiki by Mutable.ai

models

Auto-generated from tensorflow/models by Mutable.ai Auto Wiki

models
GitHub Repository
Developertensorflow
Written inJupyter Notebook
Stars76k
Watchers 2.8k
Created2016-02-05
Last updated2023-12-25
LicenseOther
Repositorytensorflow/models
Auto Wiki
Generated at2023-12-25
Generated fromCommit e923b8
Version0.0.4

The TensorFlow Models repository provides reference implementations of state-of-the-art machine learning models across multiple domains like computer vision, natural language processing, audio processing, and recommendation systems. It demonstrates best practices for modeling with TensorFlow and makes it easy to leverage TensorFlow's high-level APIs for research and development.

The repository is organized into official and research sections. The official section contains standardized implementations of popular models that are well-tested, optimized for performance, and easy to use. For example, it includes computer vision models for tasks like image classification, object detection, and segmentation built on top of common backbones like ResNet and EfficientNet. These leverage high-level Keras APIs and components like the Task class to encapsulate domain-specific workflows.

The research section provides implementations of published machine learning papers across domains. These showcase how TensorFlow can be used to recreate state-of-the-art techniques described in research publications. For example, it includes natural language processing models like BERT and computer vision models like Mask R-CNN.

A key design choice is the use of common abstractions and components across models in the official section. The Task class provides a unified interface for different machine learning problems. The Model class abstracts model building using a registry. Input pipelines leverage the flexible InputReader class. Optimization and distributed training are configured via common utilities. This enables rapid experimentation by composing reusable building blocks.

The implementations demonstrate how to leverage TensorFlow for machine learning research and development. The documentation and examples lower the barrier for users to utilize state-of-the-art models. The repository serves as a reference for best practices and provides a starting point for developing new models.

Official Models

References: official

The official directory contains standardized implementations of popular machine learning models in TensorFlow. It provides a comprehensive toolkit that is well-tested, optimized for performance, and easy to understand. The models cover a wide range of domains including computer vision, natural language processing, recommendation systems, and more.

Some key aspects of the implementation include:

  • Modular and reusable components. Common utilities like activations, losses, hyperparameters, distributed training are defined in subdirectories like /modeling, /common, /utils to be easily composed into different model types.

  • Standardized interfaces. Classes like Task and factories abstract away details of model building, data loading, and training loops. Tasks inherit from Task and implement domain-specific logic while reusing common functionality.

  • Flexible and extensible design. New models can be added by defining configurations as Python classes and registering factories. The registry pattern decouples implementations from usage.

  • Comprehensive examples. Directories like /projects provide fully-worked examples that demonstrate end-to-end usage of the framework on tasks like image classification.

  • Thorough unit testing. Tests throughout the codebase validate functionality is implemented correctly and interfaces work as intended across different scenarios.

Some important implementation details:

  • The Task class in /core/base_task.py defines the core interface for machine learning problems. Subclasses like ImageClassificationTask focus on domain logic while reusing common methods.

  • Model building is abstracted using factories, such as factory.build_classification_model() which constructs image classifiers based on configuration.

  • Configurations are defined as Python dataclasses for easy initialization, validation, and serialization. Classes like ExperimentConfig combine all necessary pieces.

  • The Trainer class in /core/base_trainer.py handles common training loop functionality so tasks focus on models, losses, and metrics. It supports both synchronous and asynchronous distributed training.

Vision

References: official/vision

The …/vision directory contains implementations of common computer vision models and components used to build models for tasks like image classification, object detection, instance segmentation, semantic segmentation, and video classification. The main functionality is organized into sub-directories:

  • …/backbones contains implementations of popular backbone networks like ResNet, MobileNet, Vision Transformer that can be used as feature extractors. The EfficientNet class implements specifications-driven backbone construction by decoding block specifications into BlockSpec objects and stacking the block modules accordingly.

  • …/decoders contains decoder networks used in vision models for tasks like semantic segmentation, including FPN, ASPP, and NASFPN. The FPN class implements Feature Pyramid Networks by adding lateral connections between backbone feature maps and constructing the top-down pathway to build a pyramid of outputs at different scales.

  • …/heads contains common head architectures for tasks like classification, detection, and segmentation including RetinaNetHead, RPNHead, DetectionHead. The RetinaNetHead class implements a common architecture for classification and box prediction towers that can be shared across levels by taking multilevel features as input and applying the towers separately to each level.

  • …/models contains implementations of complete vision models that combine components, like ClassificationModel, MaskRCNNModel, RetinaNetModel. The MaskRCNNModel class encapsulates the end-to-end Mask R-CNN algorithm by initializing components like the backbone, RPNHead, DetectionHead, ROIAligner and implementing its __call__() method to run the full pipeline.

The …/tasks directory defines several core task classes that provide a standardized interface and implementation for their domains. For example, the ImageClassificationTask class handles building and running image classification models by using the factory.build_classification_model() function to construct the model architecture based on the configuration defined in classes like BackboneConfig. The build_inputs() method leverages input readers and parsers like the Parser class which implements important preprocessing logic during parsing. Losses are defined based on the task using functions from other modules like focal_loss.FocalLoss(). The train_step() and validation_step() methods implement the core training/evaluation logic.

The …/data directory contains utilities for loading, preprocessing, and converting image datasets. It provides functions for common tasks like loading datasets via tfds.load() in /__init__.py, preprocessing images with normalize_image(), and converting datasets to TFRecords using create_coco_tf_record.py which handles parsing COCO annotations and images then encoding each example as a TfExample using TfExampleBuilder. The fake_feature_generator.py file contains utilities for generating fake feature data like images and boxes for testing purposes.

NLP

References: official/nlp

The core functionality provided by the TensorFlow NLP library for standardized natural language processing models is contained in the …/tasks directory. This directory contains implementations of common NLP tasks that can be used to easily evaluate new models. Some of the key tasks included are:

  • Question answering using the QuestionAnsweringTask class defined in …/question_answering.py. This handles building models, defining losses, loading SQuAD data, computing metrics like F1, EM scores, and running training/evaluation.

  • Sequence tagging using the TaggingTask class defined in …/tagging.py. This implements common named entity recognition and part-of-speech tagging problems with models like BERT.

  • Text classification and regression with the SentencePredictionTask class defined in …/sentence_prediction.py. This supports tasks like sentiment analysis, natural language inference and regression using models like BERT.

Each task subclass implements methods like:

This provides a consistent way to implement new NLP tasks with common methods and without needing to define lower level details.

Recommendation

References: official/recommendation

The …/recommendation directory contains implementations of various neural recommendation models. This includes implementations of Neural Collaborative Filtering (NCF) and Neural Matrix Factorization (NeuMF) models.

The core NCF framework is defined using both Keras and Estimator APIs in a standardized way. Key functionality includes:

  • Input pipelines in …/ncf_input_pipeline.py handle creating training and evaluation datasets from data producers or files. These pipelines preprocess data into a standardized format.

  • The NCF model architecture is defined using Keras layers in …/ncf_keras_main.py. This file also contains training and evaluation loops.

  • The NeuMF model architecture fuses generalized matrix factorization (GMF) and multi-layer perceptron (MLP) components. It is defined in …/neumf_model.py.

  • Common utilities for statistics, sampling, asynchronous data processing are provided in files like …/stat_utils.py, …/popen_helper.py.

  • Model training is demonstrated via the shell script at …/run.sh, which executes preprocessing, training and evaluation scripts on the MovieLens dataset.

  • The BaseDataConstructor class in …/data_pipeline.py handles asynchronously constructing training and evaluation batches.

  • The FauxPool class in …/popen_helper.py runs tasks synchronously for deterministic testing.

  • The run_ncf() function in …/ncf_keras_main.py performs end-to-end training and evaluation.

Core

References: official/core

The …/core directory contains core functionality that is reused across models in the library. This includes components like Tasks, Trainers, Inputs, and configuration definitions.

Tasks define the core machine learning problem and components. The abstract base class Task in …/base_task.py provides a common interface and infrastructure. It defines required methods to build inputs, losses, metrics, and training/evaluation steps. Concrete subclasses implement these methods to define their specific problem.

Trainers handle common training and evaluation loops. The Trainer class in …/base_trainer.py inherits from _AsyncTrainer to support both synchronous and asynchronous distributed training. It initializes important properties like the model, optimizer, and metrics. Functions like train_step() run the core training logic while updating metrics. This provides a reusable way to handle common training tasks.

Inputs are handled by the InputReader class in …/input_reader.py. It provides a configurable pipeline for input processing. The pipeline supports reading from files or TensorFlow Datasets via _read_data_source(). Data then goes through preprocessing steps like decoding, parsing, filtering and batching via configurable callbacks before being returned. This makes the input flexible to customize for different input types and domains.

Configuration definitions are provided in …/config_definitions.py. Classes like DataConfig, RuntimeConfig, and ExperimentConfig define schemas for common hyperparameters using the Schema library. This ensures configurations have a consistent structure across experiments.

Modeling

References: official/modeling

The …/modeling directory contains Python files and sub-directories that define core abstractions and utilities for building, training, and evaluating TensorFlow models. This includes reusable modeling components like activations, optimizers, and other utilities.

Some key functionality includes:

  • The …/activations subdirectory contains implementations of common activation functions used in neural networks. Activations like gelu(), mish(), relu6(), and hard_swish() are defined as TensorFlow operations and registered as Keras serializables, making them compatible with TensorFlow and Keras models.

  • The …/hyperparams subdirectory defines classes like Config and ParamsDict for representing hyperparameters in a standardized way. It supports functionality like initializing from files, overriding values, and validating against constraints.

  • The …/optimization subdirectory implements various optimization algorithms and learning rate schedules through classes like Optimizer. It supports constructing these programmatically from configuration objects. This includes optimizers like SGD, Adam, LAMB, and learning rate schedules.

Now we will discuss some important classes and functions in more detail:

The Config class defined in …/base_config.py provides a standardized way to define and override experiment configurations.

The ParamsDict class in …/params_dict.py acts as a container for hyperparameters organized in a nested dictionary structure.

The OptimizerFactory class in …/optimizer_factory.py handles constructing optimizers and learning rates from configuration objects. The build_optimizer() method constructs the actual TensorFlow optimizer object by mapping the configuration type to the appropriate class implementation.

The activation functions are implemented as simple TensorFlow operations and registered as Keras serializables. The gelu() implementation leverages the existing Keras GELU operation, while mish() directly calculates the formula from first principles. This makes the activations compatible with TensorFlow and Keras models and focuses on efficiency.

Utilities

References: official/utils

The common utilities used across models are contained in several important directories and files. The …/utils directory contains a collection of Python modules that provide common utilities for tasks like testing, documentation generation, flag configuration, and more.

Within …/utils, the /testing subdirectory implements important utilities for testing models. It provides the TestCase class and ImageCoder class which are helpful building blocks for tests. The TestCase class centralizes common testing assertions and the ImageCoder handles encoding and decoding images. The integration.py file contains the run_synthetic() function which allows running a full model training loop from tests using synthetic data. This is useful for testing for syntax errors without needing real data. The mock_task.py file defines the MockTask class for testing purposes, which implements the base Task interface and can generate dummy data and handle logging/metrics generation.

The /flags subdirectory contains utilities for defining and working with flags in a standardized way. The core.py file acts as the main coordination point, exposing functions like define_base() through a standardized API. This allows flags to be consistently configured across many models. Files like _base.py, _device.py define common flags, while validation functions like the one in _performance.py ensure valid flag combinations. The flags_test.py file systematically tests the flags functionality.

The /docs subdirectory contains utilities for programmatically generating API documentation with tensorflow_docs. The build_tfm_api_docs.py and build_orbit_api_docs.py scripts demonstrate how to generate docs by parsing models code and configurations.

The /misc subdirectory provides additional model utilities helpful for tasks like logging metrics with log(), generating synthetic data with generate_synthetic_data(), and handling callbacks and timing data with classes like TimeHistory.

Datasets

References: models

The core functionality for loading and preprocessing datasets is contained in several important locations in the codebase.

The …/datasets directory contains utilities for loading common computer vision and natural language processing datasets into TensorFlow records and example protos. The Dataset class provides a standardized interface for defining datasets across different domains. It handles loading samples from data sources via _read_data_source() and preprocessing them using functions like _parse_function(). This class presents a consistent way to access different datasets.

Within specific model implementations, datasets are often loaded and preprocessed using classes tailored to those models and domains. For example, in the DeepLab semantic segmentation implementation, the Dataset class in …/data_generator.py handles loading segmentation datasets. Its _parse_function() method applies common preprocessing steps to images and labels using functions like input_preprocess.preprocess_image_and_label(). This class encapsulates all the logic for that model's datasets in one place.

Some key utilities used across implementations for common tasks include:

Now I will discuss some important classes and functions in more detail:

The Dataset class defined in …/__init__.py provides a standardized way to access different datasets. Its _read_data_source() method handles loading samples from files or TensorFlow Datasets. _parse_function() decodes and preprocesses individual samples, applying functions from the preprocessing_factory which retrieves preprocessing functions by name. This class presents a clean, consistent interface for datasets across the library.

The pick_dataset_fn function selects the appropriate data loading based on file extensions.

The Dataset class defined in …/data_generator.py loads segmentation datasets for DeepLab. Its _parse_function() method decodes TFExamples into images and labels, and applies preprocess_image_and_label() which performs common steps such as cropping, normalization, and one-hot encoding. This encapsulates all dataset logic for that model in one place.

Pip Package

References: official/pip_package

The …/pip_package directory contains code for packaging the TensorFlow models code into a pip-installable Python package. This allows users to easily install and use the latest TensorFlow models with the pip install command.

The main file is …/setup.py, which defines the configuration for building and distributing the package. It specifies important metadata like the package name and version. The setup() function handles configuring the overall package. This function takes keyword arguments that define important aspects of the package like the name, version, description, dependencies, and packages to include.

The dependencies and their versions are defined by parsing a requirements.txt file using the _get_requirements() function. This file lists each dependency on a separate line along with the required version number or constraint. If packaging a "nightly" build, the version string is modified to include the date, and nightly versions of TensorFlow and TensorFlow Text will be used rather than stable versions.

Legacy Models

References: official/legacy

The …/legacy directory contains model implementations that are considered deprecated or legacy. These models are kept for benchmarking purposes only, as documented in the …/README.md file.

The main subdirectories are:

  • …/albert: Contains the implementation of the ALBERT model.

  • …/bert: Implements BERT models for pre-training, question answering, and classification tasks.

  • …/detection: Includes object detection models like RetinaNet, Mask R-CNN, Faster R-CNN and SSD.

  • …/image_classification: Contains implementations of models like ResNet, VGG and EfficientNet for image classification.

  • …/transformer: Provides the Transformer model for machine translation.

  • …/xlnet: Implements the XLNet model for pre-training and downstream tasks.

Some key aspects of the implementation:

The DetectionModel base class defined in …/modeling plays a key role. It defines the core interface for detection models through abstract methods like build_outputs(). Subclasses like FasterRCNN must implement these to define model architectures while inheriting common functionality.

The FasterRCNN class assembles a backbone, RPN and ROI heads using functions from …/factory.py. This implements the end-to-end Faster R-CNN model pipeline.

The ResnetRunnable class is also important, as it inherits from Orbit base classes like StandardTrainer to implement distributed training and evaluation loops via methods like train_step(). This allows easily training ResNet models.

Projects

References: official/projects

End-to-end example projects demonstrating model usage are implemented in subdirectories under …/projects. Some key examples include:

The …/text_classification_example directory contains a tutorial project for building and training a text classification model using TF-NLP. It demonstrates defining a Task subclass and key methods like build_inputs(), build_losses(), and initializing from a pretrained checkpoint. The ClassificationExampleTask class implements an end-to-end text classification task by constructing the BERT classifier model, defining the loss, and loading data. The ClassificationDataLoader class handles loading and preprocessing datasets for the task. The train.py script focuses on configuration parsing and running the general training loop logic.

The …/mtop directory contains an implementation of MTOP, a multi-task text classification model that can make predictions for all tasks with a single forward pass. The main subdirectories describe the MTOP model architecture and performance through a README file. It introduces a new multi-task news dataset and achieves superior performance compared to single-task models on news classification tasks and the GLUE benchmark.

The …/token_dropping directory implements the token dropping technique for efficiently pretraining BERT models. It defines the TokenDropBertEncoder class to implement a BERT encoder that supports token dropping. The TokenDropMaskedLMTask class inherits from MaskedLMTask and implements the masked language modeling training and evaluation loops with support for token dropping. The experiment_configs.py file registers the 'token_drop_bert/pretraining' experiment configuration defining the overall pretraining setup. The train.py file contains the main training loop.

Research Models

References: research

The research directory contains implementations of published machine learning research models across domains like computer vision, natural language processing, audio processing, and reinforcement learning. Each model implementation is contained within its own subdirectory under research, with files organized based on the domain and task.

Common utilities that support the research model implementations are provided in subdirectories like …/slim. The …/slim directory contains modules for tasks like loading and preprocessing datasets into TensorFlow records, defining common neural network layers and architectures, utilities for training loops, evaluation, preprocessing, deployment, and more. These utilities provide standardized components that make it easier to implement new models.

Some key implementation subdirectories and their purposes are:

  • …/deeplab: Contains an implementation of DeepLab models for semantic image segmentation using deep convolutional networks. It provides core layers, operations, utilities for building segmentation models, functionality for loading segmentation datasets, and modules for evaluating segmentation predictions.

  • …/attention_ocr: Implements an attention-based optical character recognition model called "Show, Attend and Read" that directly predicts character sequences from images using an attention mechanism. It contains modules for building the model architecture, training the model end-to-end, and evaluating model quality on test data.

  • …/audioset: Implements audio classification models like VGGish and YAMNet that take raw audio as input. VGGish contains code defining the pre-trained VGGish CNN architecture and preprocessing audio into log mel spectrogram examples. YAMNet defines the core YAMNet CNN architecture in the yamnet function and handles feature extraction from waveforms using waveform_to_log_mel_spectrogram_patches.

  • …/pcl_rl: Implements off-policy actor-critic and Trust Region Policy Optimization algorithms like PCL and TRPO for reinforcement learning tasks. The core training and algorithm logic is contained in the Trainer class defined in …/trainer.py. Policies and value functions are represented by classes like Policy and Baseline, while objectives are defined by classes such as PCL.

The Trainer class orchestrates the overall training process by initializing components like the EnvSpec, Model, Controller, objectives, policies, and baselines. It then runs the main training loop, periodically calling the Controller's train() method to collect batches and optimize the Model using the specified objective. The Controller interacts with environments through GymWrapper to sample episodes and store transitions in replay buffers. Objectives compute necessary quantities like the discounted future sum via discounted_future_sum() to implement the off-policy algorithm loss. Policies represent the neural network, outputting distributions and log probabilities for the controller to sample from during training.

Computer Vision

References: research/deeplab, research/delf, research/lstm_object_detection, research/vid2depth, research/object_detection

This section covers implementations of computer vision research papers in the repository. Several important directories contain code relevant to computer vision:

The …/deeplab directory contains an implementation of DeepLab models for semantic image segmentation. DeepLab uses deep convolutional neural networks to perform dense predictions on input images to classify each pixel into a semantic category. Key aspects include:

  • The …/core directory contains core functionality used in DeepLab models, including common layers, operations, and utilities needed for building segmentation architectures. It implements layers like Conv2D and operations such as separable convolutions.

  • The DensePredictionCell class in …/dense_prediction_cell.py acts as a modular "layer" to build segmentation models from configurations. Its build_cell() method iterates through operations like slim.conv2d() to build feature maps.

  • The NASBaseCell class in …/nas_cell.py implements the core recurrent logic for NAS. When called, it applies sequences of operations like separable_conv2d_same() to inputs and previous hidden states.

  • Models are defined in …/model.py using functions like multi_scale_logits() that generate outputs at multiple scales from a backbone feature extractor.

The …/evaluation directory contains code for evaluating semantic and panoptic segmentation models. The PanopticQuality class computes the Panoptic Quality metric by comparing predicted and ground truth category and instance labels.

The …/delf directory contains an implementation of Deep Local Features (DELF) for extracting local features from images. The FeatureExtractor class is used for extracting local image features using a DelfConfig. Features are aggregated using feature_aggregation_extractor.py and similarity scores computed for retrieval tasks.

Natural Language Processing

References: research/adversarial_text, research/cvt_text, research/seq_flow_lite

This section covers implementations of natural language processing research papers in the codebase. Many important NLP tasks are addressed through reusable modules and classes.

Key functionality includes:

Reinforcement Learning

References: research/efficient-hrl, research/pcl_rl

This section covers implementations of reinforcement learning research papers in the codebase. The codebase contains implementations of several important reinforcement learning algorithms and techniques.

The …/pcl_rl directory contains implementations of reinforcement learning algorithms like REINFORCE, PCL, and TRPO. These algorithms are implemented through classes like Reinforce and PCL in the …/objective.py file.

Policies and value functions are represented by the Policy class in …/policy.py and Baseline class in …/baseline.py respectively. The Policy class implements both feedforward and recurrent neural network policies. It defines methods like sample_action() to select actions and multi_step() to efficiently process full episodes for recurrent policies.

The core training and algorithm implementation is handled by the Trainer class defined in …/trainer.py. It initializes the environment, model, controller, and other components. It then runs the main training loop, collecting batches of data with the controller and optimizing the model.

Optimization algorithms for policy and value training are defined in classes like LbfgsOptimization and GradOptimization in …/optimizers.py. LbfgsOptimization implements L-BFGS optimization using both TensorFlow and SciPy for efficient optimization.

Reinforcement learning objectives are implemented through classes that define the necessary quantities for each algorithm. For example, the PCL class in …/objective.py computes the PCL loss by prepending data and calculating discounted sums with discounted_future_sum().

Replay buffers for off-policy algorithms are represented by the ReplayBuffer and PrioritizedReplayBuffer classes in …/replay_buffer.py. These implement experience replay by storing and sampling transitions from a buffer.

The Controller class in …/controller.py handles interacting with environments via the GymWrapper to collect trajectories and compute advantages for training the model. It stores transitions in the replay buffer and retrieves batches to optimize the model.

Expert trajectories for imitation learning can be generated by functions in …/expert_paths.py. These simulate algorithmic tasks to produce ground truth optimal behavior that can be used for training.

Recommendation Systems

References: research/cognitive_planning, research/rebar

The …/preprocessing directory contains implementations of recommendation system research papers. It provides functionality for loading and preprocessing multi-modal data.

The DataProcessor class in …/__init__.py handles loading and preprocessing data. It has a load_data() method that reads data from files and a preprocess() method that encodes features.

The environment is represented by the TaskEnvInterface defined in …/task_env.py. This interface specifies requirements like resetting, observing state, and mapping between graph vertices.

Policies can be trained on tasks generated from this interface. The …/tasks.py file defines the Task interface and classes for generating state sequences.

The …/policies.py file contains policy classes like LSTMPolicy for implementing recurrent agents. These policies embed states and output action probabilities.

Training is done end-to-end using code in …/train_supervised_active_vision.py. It initializes the environment, task, and policy before training the policy on state sequences in supervised manner.

Key aspects include loading multi-modal data, defining reasoning tasks, and implementing different types of reasoning agents using models like LSTMs.

Community Models

References: community

The community directory curates external TensorFlow model implementations from the open source community. It provides a central catalog and documentation for popular machine learning models across various domains.

The main file is the …/README.md which organizes models into sections for computer vision, natural language processing, and recommendation systems. It contains Markdown tables that list details about each model, including the associated paper, key features of the implementation, and maintainer of the GitHub repository where the code is hosted. Features covered include things like INT8/FP32 training and inference, automatic mixed precision, multi-GPU training with Horovod, and TensorRT support.

No code implementations are contained directly within the community directory. The purpose is to provide a curated index and documentation for users to easily find high-quality, well-maintained external TensorFlow model implementations on GitHub. This lowers the barrier to exploring state-of-the-art models. The README.md serves as the central access point, utilizing Markdown tables and links to provide information about each implementation while directing users to the external code repositories.

TensorFlow Models

References: tensorflow_models

The tensorflow_models directory provides a unified API for using models across different domains like computer vision, natural language processing, and more. It serves as the main entry point for accessing models.

The main functionality is exposed through three main modules:

  • The nlp module, defined in …/nlp, provides APIs for natural language processing tasks. It defines classes like Task for different problems and Encoder for pretrained models.

  • The vision module, defined in …/vision, exposes computer vision functionality. It contains classes for common tasks and functions for building vision models.

  • Common utilities are defined in the official module. Functionality like optimization, hyperparameters, and distributed training are imported from here.

The …/__init__.py file consolidates these modules. It imports and re-exports names from nlp, vision, and official under the parent tensorflow_models namespace. This provides a unified public interface.

Unit tests in …/tensorflow_models_test.py validate key classes and functions can be correctly imported and used. Tests are defined as methods on the TensorflowModelsTest class which inherits from tf.test.TestCase for test utilities.

NLP

References: tensorflow_models/nlp

This section covers the natural language processing functionality provided by the TensorFlow Models library. The core NLP code is contained in …/nlp, which provides a single interface for all NLP tasks, models, and functionality.

The /__init__.py file plays an important role by importing and re-exporting NLP components from other locations. It makes key items like tasks modules available in the nlp namespace. This lowers the barrier to using popular NLP techniques with TensorFlow Models.

Some important classes include:

  • Tasks classes define different NLP problems and provide methods for loading/preprocessing data.

  • Encoders contain models like BERT that extract features from text.

  • Modeling implements base model classes and layers/heads that can be composed to build custom models.

The /tasks/ subdirectory handles loading and preprocessing data for different tasks.

The /encoders/ subdirectory implements encoder models.

The /modeling/ subdirectory allows composing models with layers and heads.

Vision

References: tensorflow_models/vision

The …/vision directory contains implementations of common computer vision tasks and models using TensorFlow. The main functionality is split across several submodules:

  • The configs module defines important configuration classes for vision models. These classes standardize hyperparameters and other settings for models.

  • The modeling module contains classes and functions for building vision models. The base class defines the common interface for vision models, including methods like call() for running predictions. The build_model() function constructs the actual Keras model graph based on the configuration. It handles common preprocessing like resizing and normalizing inputs.

  • The tasks module defines common computer vision tasks with classes to load and preprocess image data and evaluate models. The ImageClassificationTask handles loading image datasets and applying preprocessing. It can then evaluate models by calculating metrics like accuracy on a test set. Datasets for tasks are stored in TensorFlow Datasets format for easy loading.

Utilities

References: models

The core utilities for common tasks like testing, documentation generation, configuration, callbacks, and synthetic data generation are provided in the …/utils directory. This directory contains Python modules that implement reusable utilities for tasks like training and evaluating TensorFlow models. Some key functionality includes:

Testing utilities are provided in …/testing. This includes the TestCase class that is subclassed for writing unit tests. It provides common utilities like random_image() that generate dummy data useful for tests.

Documentation generation utilities are located in …/docs. These utilize tensorflow_docs to generate API references and documentation from code docstrings in a standardized format.

Flag definition utilities are contained in …/flags. The core.py file acts as the main coordination point, exposing functions like define_base() through a standardized API. This provides a consistent way to configure models with flags across the library.

Callbacks and training utilities are defined in …/keras_utils.py. This includes the TimeHistory callback class that can be used to log elapsed time during training.

Synthetic data generation functions are provided in …/model_helpers.py. These are useful for generating dummy inputs to test or benchmark models without a dependency on real datasets.

Tests

References: tensorflow_models/tensorflow_models_test.py

The …/tensorflow_models_test.py file contains unit tests for core functionality in the tensorflow_models library. The TensorflowModelsTest class defines several test methods to validate important components.

The testVisionImport() method tests the vision module by instantiating a SqueezeExcitation layer and importing the Losses config class.

The testNLPImport() method similarly tests the nlp module by instantiating a TransformerEncoderBlock layer and importing the TaggingTask class.

The testCommonImports() method ensures the common utilities work by importing and instantiating the Config class and a LinearWarmup schedule.

Each test method leverages TensorFlow's testing utilities by inheriting from tf.test.TestCase. This allows methods to validate classes and functions are imported and work as expected. Running all tests ensures the overall stability of the library.

Utilities

References: official/utils, official/common

The common utilities used across models are primarily contained in the …/__init__.py and …/misc directories.

The …/__init__.py file contains classes that standardize important aspects of model building. The DataConfig class defines attributes like image_size and batch_size to ensure consistent input data preprocessing. The ModelBuilder class provides a standardized interface for constructing model architectures using a registry. The OptimizerConfig and LearningRateSchedule classes encapsulate hyperparameters for training models in a configurable way. The Validation class handles evaluating models on validation/test data.

In …/__init__.py, the log(), metric(), and metric_obj() functions handle logging metrics to TensorBoard. The log() function logs a scalar value, while metric() logs a named metric value. metric_obj() returns a Metric object that can be used to aggregate metric values over multiple steps or epochs before logging.

The …/keras_utils.py file contains the TimeHistory Keras callback class, which logs timing data like steps/second during training. In its on_batch_end() method, it calculates these metrics and logs them to a SummaryWriter if enough steps have elapsed. The …/model_helpers.py file provides generally useful model utilities like the past_stop_threshold() function, which checks if a model should stop training based on an evaluation metric threshold.

Testing

References: official/utils/testing, official/utils/testing/scripts

The core utilities for testing models and components are contained in …/testing. This directory provides common building blocks and abstractions that make testing TensorFlow models easier, cleaner, and less dependent on external data sources.

Some key functionality includes the TestCase class, which subclasses tf.test.TestCase and adds additional helper methods for testing like assertNear(). This centralizes common testing assertions and avoids duplicating logic in tests.

The ImageCoder class handles encoding and decoding of images for testing purposes. It can encode images to TF example protos and decode them back, removing the need for tests to handle image serialization formats directly.

Functions like random_image() make it easy to generate synthetic test data without requiring real images. This allows tests to run quickly without external data dependencies.

The MockTask class implements the base Task class needed for models. It defines a simple Keras sequential model in MockModel and generates dummy data to track loss for validation. MockTask is useful for testing purposes as it implements the core task interface without dependencies on real data.

The integration.py module contains the run_synthetic() function, which allows running a full model training loop directly from tests. It sets up flags, runs the model main using synthetic data generated on the fly, and cleans up temporary directories after, providing a minimal way to test for syntax errors in the training loop implementation.

Unit tests for individual components are contained under subdirectories like …/scripts. The presubmit.sh script implements the py_test() function which runs each Python test file sequentially, collecting exit codes to return the overall test result. This avoids race conditions when running tests in parallel.

Documentation

References: official/utils/docs

The …/__init__.py file contains helper functions for extracting information from model configurations.

The build_tfm_api_docs.py script in …/docs generates API documentation for TensorFlow Models. It uses the DocGenerator class from tensorflow_docs to parse TFM code and generate Markdown files.

DocGenerator drives the documentation building. These scripts programmatically document TFM models and APIs.

Configuration

References: official/utils/flags

The …/flags directory contains utilities for configuring models through flags (also known as hyperparameters).

Key functionality includes:

  • The define_base() function defines common base flags like the data and model directories, batch size, and number of epochs. These flags ensure consistency for I/O and training hyperparameters.

  • Functions like define_device() and define_distribution() allow configuring flags for hardware settings and distributed execution. This includes flags for the TPU name, number of shards, and worker hosts list.

  • Flag values can be validated using validator functions decorated with @flags.validator. An example is the _check_benchmark_log_dir validator in _benchmark.py which checks the log directory is set appropriately.

  • Default flag values are set. This allows reasonable defaults to be specified while still allowing overrides.

  • The flags_test.py module systematically tests the flags functionality to prevent regressions as the API evolves. It validates behaviors like parsing, default values, and data types.

Distributed Training

References: official/common/distribute_utils.py

The …/distribute_utils.py file provides utilities for running models in a distributed setting using different TensorFlow distribution strategies. It handles configuring the distribution strategy based on parameters like number of GPUs or TPU address.

The _mirrored_cross_device_ops() function initializes a CrossDeviceOps object for MirroredStrategy based on the all-reduce algorithm and number of packs. It returns either a tf.distribute.NcclAllReduce or tf.distribute.HierarchicalCopyAllReduce object configured for collective operations.

The get_distribution_strategy() function takes strategy name and other parameters to instantiate and return the appropriate distribution strategy class. It passes the initialized CrossDeviceOps to MirroredStrategy.

For multi-worker training, configure_cluster() sets the TF_CONFIG environment variable to configure TensorFlow jobs by parsing worker hosts and task index. This allows running distributed training with MultiWorkerMirroredStrategy.

The tpu_initialize() function initializes TPU systems by connecting to a TPU cluster resolver from the TPU address. It configures TensorFlow to connect to this cluster.

Callbacks

References: official/utils/misc/keras_utils.py

The Keras callbacks defined in keras_utils.py help with monitoring and checkpointing during model training. The TimeHistory callback logs timing statistics like steps/second and epoch runtimes to track training progress. It calculates these metrics periodically and logs them via TensorFlow summaries. The SimpleCheckpoint callback saves checkpoints by calling save() on a tf.train.CheckpointManager object. It saves after each epoch for easy recovery of trained models.

The TimeHistory class tracks step and batch timing data during training. It overrides the on_batch_end() and on_epoch_end() methods to calculate metrics like steps/second when enough batches have processed. These metrics are either logged directly or added to TensorFlow summaries if a SummaryWriter is provided. on_epoch_begin() and on_train_end() methods track epoch start/end times.

The SimpleCheckpoint class saves model checkpoints by calling save() on a CheckpointManager passed into its constructor. It overrides on_epoch_end() to trigger a save operation at the end of each epoch for easy checkpointing.

The set_session_config() function configures the TensorFlow session, enabling optimizations like XLA if specified. The set_gpu_thread_mode_and_count() function sets environment variables to control GPU thread counts and the TensorFlow dataset thread pool size, balancing runtime needs.

Synthetic Data

References: official/utils/misc/model_helpers.py

The file …/model_helpers.py contains utilities for generating synthetic data using TensorFlow. The generate_synthetic_data() function takes in input and label shapes and constant values to generate a repeating synthetic dataset from constant tensors. This provides a simple way to test models on sample data without needing real datasets.

The generate_synthetic_data() function first uses tf.data.Dataset.from_tensors() to create tensors from the provided input and label shapes and constant values. It then chains a call to .repeat() to create an infinite repeating dataset from these constant tensors. This dataset can be iterated over for testing or to quickly generate synthetic batches, without requiring any real data processing or I/O. The use of TensorFlow data APIs like Dataset.from_tensors() and .repeat() allows the generated data to integrate cleanly into models and training loops expecting Dataset objects.

Datasets

References: research/slim/datasets, research/vid2depth/dataset

This code provides functionality for loading, preprocessing, and converting datasets for use in machine learning models. Datasets are standardized into common formats that can be easily accessed by models through standardized interfaces. Raw datasets are preprocessed into standardized formats that are efficient for training models.

The …/datasets directory contains utilities for loading and preprocessing many common computer vision and NLP datasets. The core abstraction is the Dataset class defined in /__init__.py, which provides a base interface for defining datasets. Child classes like CIFAR10 in …/cifar10.py implement methods for specific datasets to handle tasks like downloading, parsing, and preprocessing data. The get_split() method on each dataset class returns a Dataset namedtuple containing instructions for loading a split of data for processing.

The …/dataset directory contains functionality for loading video-to-depth datasets. Classes like KittiRaw() in …/dataset_loader.py abstract accessing different raw datasets and allow generating preprocessed data in a source-agnostic way. The Dataset class in /__init__.py provides a common interface for loading preprocessed datasets.

Files like …/download_and_convert_cifar10.py show end-to-end workflows to download raw datasets and preprocess them into standardized TFRecord format. Utilities in …/dataset_utils.py help with tasks like populating TFExample protos using functions like image_to_tfexample().

The gen_data.py file in …/dataset contains the Dataset class which handles preprocessing samples from different datasets into a standardized format and saving efficiently to disk. It uses multiprocessing via Python's Pool to maximize throughput when generating examples.

Vision

References: models

The standardized computer vision models implemented in the …/vision directory provide a comprehensive and reusable toolkit for computer vision research in TensorFlow. The core functionality includes:

Model configurations are defined as Python dataclasses, providing a consistent interface for selecting hyperparameters like network architectures. The Backbone class in …/backbones.py acts as a container class that can hold different backbone configurations using a OneOfConfig.

Datasets are loaded and preprocessed using utilities in …/data and …/dataloaders. The Parser class in …/classification_input.py handles common preprocessing steps on image data.

Vision tasks like image classification are implemented in classes such as ImageClassificationTask in …/image_classification.py. This class encapsulates building the model, defining losses, loading data, computing metrics, and running the training loop.

Model building is abstracted using factories like factory.build_classification_model() in …/factory.py, which constructs image classifiers based on configuration. This factory approach allows flexibly building different model architectures.

Losses are defined modularly, with the FocalLoss class in …/focal_loss.py handling class imbalance issues.

Evaluation metrics are implemented in classes such as COCOEvaluator in …/coco_evaluator.py for object detection.

Model export is standardized via modules like ClassificationModule in …/image_classification.py.

NLP

References: models

The …/nlp directory provides standardized implementations of popular natural language processing models for tasks like question answering, language modeling, and machine translation. It contains reusable components, implementations of well-known NLP architectures, training utilities, and documentation.

Some core functionality includes:

  • The …/data directory contains utilities for loading and preprocessing various NLP datasets into consistent TensorFlow formats. This includes tokenization, preprocessing pipelines handled by functions like BasicTokenizer, and data loaders.

  • The …/modeling directory implements common neural network layers, models, and other primitives for building transformer-based NLP architectures.

  • Tasks are implemented through classes such as QuestionAnsweringTask that encapsulate full pipelines for question answering and sequence tagging respectively.

  • Training is handled by functions in using the run_experiment() loop. This initializes distributions, loads the task and model, and runs training.

  • Configuration options are defined in classes like EncoderConfig in. These provide a consistent interface for configuring experiments.

Some important implementation details:

  • The EncoderConfig class uses the OneOfConfig pattern to allow flexible encoder configurations via a unified interface.

Recommendation

References: models

The …/recommendation directory contains standardized implementations of popular recommendation system models in TensorFlow. It provides modular and reusable implementations of neural collaborative filtering models like Neural Collaborative Filtering (NCF) and Neural Matrix Factorization (NeuMF).

The core NCF framework is defined using both Keras and Estimator APIs. Key files include:

For data preprocessing it includes:

It provides utilities for statistics, sampling, and asynchronous data processing in files like:

Model training is demonstrated via the shell script at …/run.sh. Comprehensive unit tests are in …/ncf_test.py.

The BaseDataConstructor class in …/data_pipeline.py is the base class for different data preprocessing pipelines. Its subclasses like MaterializedDataConstructor handle asynchronously constructing training and evaluation batches via methods like make_input_fn().

The FauxPool class in …/popen_helper.py mimics a multiprocessing pool but runs tasks synchronously, enabling deterministic testing of asynchronous preprocessing code.

The run_ncf() function in …/ncf_keras_main.py executes end-to-end training and evaluation, leveraging other classes like MetricLayer for metrics computation.

Core

References: models

The core components used across models in TensorFlow include tasks, trainers, inputs, and configurations. Tasks define machine learning problems and encapsulate models, losses, and training/evaluation loops for specific domains. The Task class in …/base_task.py provides a common interface for tasks by defining abstract methods like build_inputs(), build_losses(), and train_step() that child tasks implement.

Trainers handle common training and evaluation logic. The Trainer class in …/base_trainer.py initializes important properties and runs training methods like train_step() and evaluation methods like eval_step() to optimize models and track metrics using tasks. It supports both synchronous and asynchronous distributed training.

Inputs are processed using flexible input pipelines. The InputReader class in …/input_reader.py provides a configurable pipeline for input processing via _read_data_source() and preprocessing callbacks before batching examples. It can read from files or TensorFlow Datasets.

Configurations define hyperparameters. The Config class in …/__init__.py stores attributes in a standardized way for tasks like data preprocessing. The OptimizerConfig and LearningRateSchedule classes also define hyperparameters.

The Task class is central to tasks - it defines the interface via abstract methods and handles common functionality. Child tasks implement methods to define their specific problems. The Trainer handles common training/evaluation logic so tasks focus on models/losses. The InputReader class provides a modular input pipeline with callbacks for flexible preprocessing. Configurations like Config standardize hyperparameters.

Modeling

References: models

The …/modeling directory contains Python files and sub-directories that define core abstractions and utilities for building, training, and evaluating TensorFlow models. It provides reusable components like neural network layers, activation functions, and optimization algorithms that can be used to construct machine learning models in a modular and configurable way.

Some of the key functionality included under Modeling are:

  • The …/activations subdirectory contains implementations of common activation functions like gelu(), mish(), relu6(), and hard_swish(). These are defined as TensorFlow operations and registered as Keras serializables, providing a standardized interface for using different activation functions in models.

  • The …/hyperparams subdirectory defines classes like Config and ParamsDict for representing hyperparameters in a standardized way. The Config class supports functionality for initializing from files, overriding values, and validating against constraints.

  • The …/optimization subdirectory implements various optimization algorithms and learning rate schedules through classes like Optimizer. The OptimizerFactory class handles constructing optimizers and learning rates from configuration objects.

  • The …/performance.py file contains utilities for mixed precision training, including functions like configure_optimizer() and set_mixed_precision_policy(). These make it easier to optimize models in float16 while preventing underflow issues.

Some important classes and functions related to modeling include:

Utilities

References: models

The core common utilities used across models are contained in the …/common and …/utils directories. These directories provide reusable components, classes, and functions that are frequently used when building models.

The Config class defined in …/__init__.py provides a standardized way to define and override experiment configurations. It supports functionality like initializing from files, overriding values, and validating against constraints.

The functions in …/distribute_utils.py help configure distribution strategies for distributed training. The _mirrored_cross_device_ops() function configures collective operations like all-reduces for MirroredStrategy by mapping the algorithm to classes like NcclAllReduce that perform the collective communication efficiently across devices. The get_distribution_strategy() function handles mapping parameters to the correct configuration values for each strategy, enabling flexible and consistent distributed training.

The TestCase class provided in …/__init__.py includes utilities like random_image() that are useful for writing unit tests.

Testing utilities are also provided in …/testing including the TestCase class that can be subclassed for authoring tests.

Datasets

References: models

The core functionality for loading and preprocessing datasets in the TensorFlow Models codebase is contained in several key locations:

The …/datasets directory contains utilities for loading common computer vision and natural language processing datasets into TensorFlow records and example protos. The Dataset class provides a standardized interface for defining datasets across different domains. It handles loading samples from data sources like files or TensorFlow Datasets via its _read_data_source() method.

Some important datasets that are supported include ImageNet, COCO, and MNIST, which are loaded by functions in files like imagenet.py, coco.py, and mnist.py. These functions preprocess data like images by resizing and normalizing them. The standardized Dataset interface allows these datasets to be used interchangeably in different parts of the library.

The pick_dataset_fn function in …/dataset_fn.py is also important for loading datasets. It selects the appropriate TFRecordDataset class or partial function to use based on whether files are plain or compressed. This avoids duplicating code for different file types. The PossibleDatasetType union allows it to return either a class or callable in a type-safe manner.

In the computer vision domain, the …/dataloaders directory contains additional functionality. The Parser class handles common preprocessing steps on image datasets like resizing, normalization, and parsing annotations. It wraps examples for models using the standardized InputReader interface from …/core.

The DeepSpeechDataset class in …/data handles important preprocessing for speech recognition datasets. It extracts log mel filterbank features from raw audio via the AudioFeaturizer class. It then sorts examples by length and generates minibatches of fixed-size examples as TensorFlow datasets.

The load_data() function in …/datasets.py presents a clean API for loading different NLP and vision datasets into a consistent format. It handles tasks like binarizing MNIST, preprocessing text, and returning standardized train/valid/test splits. Configuration for filenames and locations is separated out into …/config.py

Pip Package

References: models

The …/pip_package directory contains code to package the TensorFlow models code into a pip-installable Python package. This allows users to easily install and use the latest TensorFlow models with the pip install command.

The core functionality is contained in the …/setup.py file. This file defines the configuration needed for building and distributing the package. It specifies important metadata like the package name and version. The setup() function handles configuring the overall package. When users run pip install, this setup.py file ensures the correct files are included in the installation.

Some key aspects of how it works:

  • The package name, version number, and other metadata are defined using standard Python packaging fields at the top of …/setup.py.

  • The setup() function handles tasks like specifying package dependencies, including necessary files in the package, and providing a descriptive long_description for display on PyPI.

  • It registers the package on PyPI using the setup() call, making it available for installation.

  • When a user runs pip install, it will download the sdist archive from PyPI containing all necessary files.

  • pip will then unpack and install the package by running the setup.py file, ensuring all files are placed in the correct location on the user's system.

So in summary, the …/setup.py file contains all the necessary configuration to build a Python package that can be installed via pip. It handles important packaging tasks like metadata, dependencies, and file inclusion to distribute the TensorFlow models code as a reusable package.

Legacy Models

References: models

The …/legacy directory contains implementations of popular machine learning models that are considered legacy or deprecated. These models are kept for benchmarking purposes only, as documented in the …/README.md file.

The main subdirectories are:

  • …/albert: Contains the implementation of the ALBERT model.

  • …/bert: Implements BERT models for pre-training, question answering, and classification tasks.

  • …/image_classification: Contains implementations of models like ResNet, VGG and EfficientNet for image classification.

  • …/transformer: Provides the Transformer model for machine translation.

  • …/xlnet: Implements the XLNet model for pre-training and downstream tasks.

Some key aspects of the implementation include:

The ResnetRunnable class is also important, as it inherits from Orbit base classes like StandardTrainer to implement distributed training and evaluation loops via methods like train_step(). This allows easily training ResNet models.

Projects

References: models

The …/projects directory contains complete example projects that demonstrate end-to-end usage of the TensorFlow Models library. These projects serve to illustrate how all the pieces like data input, model building, training, evaluation, export fit together into a cohesive workflow.

Some key example projects include:

The business logic focuses on defining Task subclasses that encapsulate all aspects of a given problem - from data input via build_inputs(), model construction with build_model(), calculating losses in build_losses(), and implementing the core training/evaluation loops by overriding methods like train_step() and validation_step(). This provides a clean, reusable interface for different tasks. The train.py scripts tie together preprocessing, model training/evaluation, checkpointing, and exporting into an end-to-end workflow.

Research Models

References: models

The research subdirectory contains implementations of published machine learning research models across various domains like computer vision, natural language processing, audio processing, and reinforcement learning. Each subdirectory focuses on a particular domain or technique, containing modular implementations of models from seminal papers in that area.

The …/deeplab directory includes code for building segmentation models using the DeepLab architecture. The DensePredictionCell class acts as a reusable "layer" to construct models from configurations. Models are defined modularly using classes like DensePredictionCell and functions from /core and /model.py. The …/object_detection directory provides a standardized framework that makes it easy to construct, train and deploy popular object detection models with TensorFlow. The DetectionModel class defines the core interface that all detection models must implement.

The …/audioset directory contains code and pre-trained models related to the AudioSet dataset. It includes implementations of models like VGGish and YAMNet for audio feature extraction and sound classification. The VGGish model is defined in vggish_slim.py, applying convolutions in a VGG-like fashion.

The …/cognitive_planning directory provides environments, tasks, policies, and utilities for training agents on embodied AI problems. The Trainer class defined in trainer.py orchestrates the overall RL training process.

For implementation details on specific models, please refer to the relevant subdirectory sections in this wiki.

Computer Vision

References: research/delf, research/lstm_object_detection, research/vid2depth, research/object_detection

The core functionality implemented in the …/object_detection directory is computer vision algorithms and models. This includes implementations of popular object detection architectures and research papers in computer vision domains like vision, NLP, audio processing, and reinforcement learning.

Some key implementations include:

The key aspects of the implementation are defining common interfaces for components like DetectionModel and AnchorGenerator, implementing popular CNN backbones, object detection models from papers, and utilities for training/evaluation. Thorough unit tests validate the algorithms.

Natural Language Processing

References: research/cvt_text, research/seq_flow_lite

This section covers implementations of natural language processing research papers in the repository. The core functionality is split across several important directories:

The …/cvt_text directory implements the Contextualized Vector Text (CVT) model, a multi-task semi-supervised model for NLP. It provides a full reference implementation of the CVT technique. The modular design allows flexibility in configuring new tasks and models for multi-task semi-supervised learning using both labeled and unlabeled data as per the CVT algorithm.

The …/seq_flow_lite directory contains implementations of sequence modeling architectures and related utilities for NLP tasks. It defines several custom ops like DenylistOp and SequenceStringProjectionOp for tasks like text processing and sequence decoding. Text can be efficiently embedded, normalized, tokenized, and distorted for augmentation. Skipgrams and subsequences can be found lazily using tries. Beam search decoding is implemented in the BeamSearch class.

The key implementation details are:

Audio Processing

References: research/deep_speech, research/lfads

The …/data directory contains utilities for loading and preprocessing audio data for training deep speech models. It implements common techniques for feature extraction from raw inputs to generate tensors that can be directly used for model training and inference.

The main classes and functions are defined in the …/featurizer.py file. The AudioFeaturizer class initializes parameters for extracting spectrogram features from audio like sample rate, window length and stride. The core functionality is in the compute_spectrogram_feature() function, which takes in audio samples, sample rate and other parameters and computes the spectrogram features using a short-time Fourier transform approach. It first extracts overlapped windows from the input samples using a stride that is smaller than the window size. It then applies a Hann window function to each window, takes the squared magnitude of the FFT, and scales the result. It returns the log of the spectrogram plus a small epsilon value to avoid issues with taking log of 0.

The DeepSpeechDataset class handles important preprocessing steps like normalizing audio features and sorting the dataset entries by audio length. It leverages the featurizers to load audio files, extract features, parse a CSV file to get tuples of audio and transcript data, and generate TensorFlow datasets via the input_fn() function by preprocessing, shuffling, batching and repeating the data.

The …/download.py file contains utilities for downloading the LibriSpeech dataset from URLs, extracting archives, converting audio formats and splitting transcripts via functions like download_and_extract() and convert_audio_and_split_transcript(). The download_and_extract function handles downloading a tarball from a URL and extracting it, using tempfile.mkstemp to download to a temporary file and tarfile.open to extract the contents to the target directory. The convert_audio_and_split_transcript function takes an input directory, uses the sox library via the Transformer class to convert FLAC to WAV format, and generates a CSV file with audio and transcript data.

Reinforcement Learning

References: research/pcl_rl

This section covers implementations of reinforcement learning research papers in the areas of computer vision, natural language processing, audio processing, recommendation systems, and more.

The core functionality is contained in the …/pcl_rl directory, which provides implementations of reinforcement learning algorithms like REINFORCE, PCL, TRPO, and UREX. These algorithms are represented by classes like Reinforce, PCL, and TRPO defined in objective.py. These classes compute the necessary losses, targets, and gradients for training RL agents.

Policies and value functions are represented by neural network models defined in policy.py and baseline.py. The Policy class implements both recurrent policies using LSTMs as well as feedforward MLP policies. It handles functions like sampling actions and computing entropy, KL divergence, and log probabilities. The Baseline class defines the architecture for value estimation.

The Model class in model.py brings together the policy, baseline, objective, and optimizer to define the full TensorFlow graph for training. It supports different objectives, network architectures, and optimization algorithms. Key methods include sample_step() for sampling and train_step() for optimization.

Interaction with environments is handled by the GymWrapper class in gym_wrapper.py. It runs multiple Gym environments simultaneously in a batched manner while tracking internal states like observations and rewards over time.

Optimization uses classes like TrustRegionOptimization defined in trust_region.py. It implements trust region policy optimization efficiently via conjugate gradients without explicit matrix operations.

Experience replay is supported via the ReplayBuffer and PrioritizedReplayBuffer classes in replay_buffer.py, which store and sample transitions for off-policy training.

The Controller class in controller.py interacts with environments via GymWrapper to collect trajectories and compute advantages, then trains the Model.

Expert demonstrations can be generated by functions in expert_paths.py that simulate algorithmic tasks to provide ground truth trajectories for imitation learning.

The Trainer class in trainer.py coordinates training by initializing all components and running the main optimization loop. It handles checkpointing models.

Recommendation Systems

References: research/rebar

The …/rebar directory contains an implementation of the REBAR algorithm applied to collaborative filtering for recommendation systems. The core REBAR algorithm is defined in the SBN class in the file …/rebar.py, which implements functions like _random_sample() and _create_baseline() for backpropagating through stochastic nodes. Training and evaluation is handled by …/rebar_train.py. Datasets are loaded using the load_data() function defined in …/datasets.py, and metrics are logged with the Logger class defined in …/logger.py. The SBN class contains the core model, loss, and optimization logic for training models using REBAR gradient estimation.

Utilities

References: models

The core utilities used across research model implementations are contained in the …/slim directory. This directory provides common utilities that are frequently reused when building models.

Some key functionality includes:

The …/datasets subdirectory contains utilities for loading and preprocessing popular computer vision and natural language processing datasets. The Dataset class provides a standardized interface for defining datasets that can then be easily loaded and iterated over.

The …/preprocessing subdirectory implements common preprocessing steps for images and text. This includes functionality for cropping, normalization, and other standard preprocessing techniques. The preprocessing_factory function retrieves the appropriate preprocessing function based on the type of data.

Utilities for tasks like configuration are provided in the …/deployment subdirectory. The DeploymentConfig class controls how models are deployed to different devices.

Community Models

References: models

The community directory curates external TensorFlow model implementations for users to reference. It provides documentation on relevant models from the literature along with links to GitHub repositories containing implementations of those models powered by TensorFlow 2.

The core functionality is contained in the …/README.md file. This file organizes the models into sections for computer vision, natural language processing, and recommendation systems using Markdown tables. Each section contains details about the models, associated papers, key features of the implementations, and maintainers for each implementation linked on GitHub.

Some important implementation details:

  • The README.md file provides the central documentation and index of external TensorFlow model implementations. It uses Markdown tables and links to organize information about different models.

  • No complex code logic or algorithms are implemented directly within the community directory. The main purpose is to surface external resources through documentation in README.md.

  • Features documented in README.md include things like quantization support, distributed training, and deployment options. This gives users an overview of each implementation's capabilities.

  • Maintainers linked include organizations providing pre-trained models like Intel and NVIDIA, as well as individual external contributors. The links allow crediting work and finding support resources.

TensorFlow Models

References: models

The tensorflow_models directory serves as the main entry point for accessing functionality from the TensorFlow Models library. It provides APIs for common natural language processing and computer vision tasks through modules in the …/nlp and …/vision subdirectories. These modules define reusable components, implementations of popular models, and utilities for tasks like preprocessing, training, and evaluation that can be easily leveraged in other projects.

The …/__init__.py file acts as the primary import hub. It collects exports from other modules under the tensorflow_models namespace, providing a unified interface. This includes functionality from nlp, vision, and official without requiring importing from each directly. No complex logic is implemented in /__init__.py - it simply re-exports names to make other modules accessible.

Some important classes and functions include:

  • The Task class defined in …/nlp provides a common interface for different NLP tasks like question answering or named entity recognition. It can be subclassed to implement new tasks.

  • The ClassificationModel class in …/vision provides a standardized way to build image classification models.

  • Preprocessing utilities for vision are defined in …/vision.

  • The model_fn function in …/tensorflow_models_test.py constructs and trains simple models for unit testing purposes.

NLP

References: models

This section covers natural language processing functionality provided by the code. The core NLP functionality is contained in the …/nlp directory. This directory contains a library for natural language processing in TensorFlow that provides reusable components, implementations of popular NLP models, training utilities, and documentation.

Some key aspects of the NLP functionality include:

  • The …/data subdirectory, which contains utilities for loading and preprocessing various NLP datasets into TensorFlow formats. These utilities support tasks like tokenization and preprocessing pipelines.

  • The …/modeling subdirectory, which contains reusable Keras layers, models, and other primitives for building transformer-based NLP architectures. It implements important components like attention layers.

  • The …/tasks subdirectory, which implements common NLP tasks such as question answering and sequence tagging through classes like QuestionAnsweringTask and TaggingTask. These classes define the full task pipeline.

Now I will discuss some important classes and functions in more detail:

The QuestionAnsweringTask class in …/question_answering.py is central to building question answering models.

The train_lib.py module contains the core multi-GPU training loop implementation. The run_experiment() function handles initializing the optimizer, running training iterations, evaluating periodically, and saving checkpoints. This provides a standardized training loop.

The BasicTokenizer class in …/tokenization.py implements basic tokenization. Its tokenize() method takes text and returns tokenized tensors. This provides a clean interface for tokenization.

Vision

References: models

The core computer vision functionality in the TensorFlow Models repository is implemented in the …/vision directory. This directory contains standardized implementations of common computer vision models and components for tasks like image classification, object detection, and semantic segmentation.

Some key aspects of the functionality in …/vision include:

  • Configurations for vision models, datasets, and experiments are defined as Python dataclasses in the /configs subdirectory. For example, the Backbone class in /configs/backbones.py defines common backbone architectures.

  • Image preprocessing utilities are provided in the /data subdirectory. Functions like normalize_image() in /data/__init__.py apply common preprocessing steps to images.

  • Input pipelines for parsing images and annotations are implemented in the /dataloaders subdirectory. The Parser class in /dataloaders/classification_input.py handles preprocessing inputs.

  • Evaluation metrics for tasks like object detection are defined in subdirectories such as /evaluation/coco_evaluator.py.

  • Example projects demonstrating end-to-end usage are located in /examples, such as the image classification starter in /examples/starter.

  • Losses functions for tasks like object detection are implemented in /losses/focal_loss.py.

  • Modeling components like backbone networks, decoders, and heads are defined in /modeling. The ClassificationModel class in /modeling/classification_model.py constructs image classifiers.

  • Operations for common computer vision tasks are implemented in /ops. The AnchorGenerator class in /ops/anchor_generator.py generates anchors for object detection.

  • Computer vision problems are encapsulated as classes inheriting from Task in /tasks. For example, the ImageClassificationTask handles training classifiers.

  • Utilities for tasks like object detection are provided in subdirectories like /utils/object_detection/box_list.py.

Now I will discuss some important classes and functions in the …/vision directory in more detail:

The Backbone class in /configs/backbones.py acts as a container that can hold the configuration of any supported backbone model using a OneOfConfig. This provides a unified interface for selecting different backbone architectures like ResNet, EfficientNet, Vision Transformer etc.

The ImageClassificationTask class in /tasks/image_classification.py is central to building image classification models. Its build_model() method uses the factory.build_classification_model() function to construct the Keras model architecture based on the configuration.

The MaskRCNNModel class in /modeling/maskrcnn_model.py encapsulates the end-to-end Mask R-CNN algorithm. It initializes components like the backbone, RPNHead, DetectionHead, ROIAligner and its __call__() method runs the full pipeline.

The AnchorGenerator base class in /ops/anchor_generator.py provides a common way to generate anchors tailored for specific model architectures. Subclasses like GridAnchorGenerator must implement _generate() to produce the anchors.

Utilities

References: models

The core utilities for common tasks like testing, documentation generation, configuration, callbacks, and synthetic data generation are provided in the …/utils directory. This directory contains Python modules that implement reusable utilities for tasks like training and evaluating TensorFlow models. Some key functionality includes:

Testing utilities are provided in …/testing. This includes the TestCase class that is subclassed for writing unit tests. It provides common utilities like random_image() that generate dummy data useful for tests.

Documentation generation utilities are located in …/docs. These utilize tensorflow_docs to generate API references and documentation from code docstrings in a standardized format.

Flag definition utilities are contained in …/flags. The core.py file acts as the main coordination point, exposing functions like define_base() through a standardized API. This provides a consistent way to configure models with flags across the library.

Callbacks and training utilities are defined in …/keras_utils.py. This includes the TimeHistory callback class that can be used to log elapsed time during training.

Synthetic data generation functions are provided in …/model_helpers.py. These are useful for generating dummy inputs to test or benchmark models without a dependency on real datasets.

Tests

References: models

The core functionality validated by unit tests in this codebase is encapsulated by the TestCase class defined in …/__init__.py. This class provides a base for writing unit tests that can be run with utils_test() to validate components.

Some key aspects of how it works:

  • TestCase inherits from tf.test.TestCase and adds TensorFlow specific assertions like assertNear() for comparing floating point values. This allows tests to leverage TensorFlow's testing utilities.

  • Tests are organized into classes inheriting from TestCase, with each test defined as a method.

  • Within test methods, the TestCase instance's assertion methods can be used to validate outputs against expected values. For example self.assertEqual(actual, expected) to check equality.

  • Additional helper classes are provided like MockTask in …/__init__.py which implements the base Task interface needed for testing model functionality without depending on a real task. This improves test isolation.

  • The utils_test() function defined in …/__init__.py acts as the main test runner. It uses TensorFlow's test.main() to discover and execute all unit tests within the utils module, validating core functionality.

Utilities

References: official/common

This section covers common utilities, tests, configurations, and components used across models in the repository. Key functionality includes:

The …/common directory contains many reusable components that standardize important aspects of model building. The ModelBuilder class provides a standardized interface for constructing model architectures from a configuration. It uses a registry to map names to functions that build architectures. The DataConfig class defines attributes for standardizing input preprocessing.

Testing utilities ensure components work reliably. The tests in …/distribute_utils_test.py validate distribution functions under different scenarios. This helps ensure correct distributed behavior.

Configuration classes encapsulate hyperparameters. The OptimizerConfig class encapsulates learning rate and optimizer parameters. LearningRateSchedule implements learning rate schedules as callable functions, providing flexibility.

The Validation class handles evaluating models on validation data. pick_dataset_fn in …/dataset_fn.py selects the appropriate loading function based on file type, avoiding duplicated code.

Functions in …/distribute_utils.py help configure distribution strategies. _mirrored_cross_device_ops() maps the all-reduce algorithm to classes like NcclAllReduce that perform efficient collective communication. get_distribution_strategy() handles mapping parameters to correct configuration values.

Global counters defined in …/streamz_counters.py allow tracking specific metrics like counts of operations. These provide visibility into model behavior during training.

Common flags defined in …/flags.py avoid issues from inconsistent flag definitions across libraries. The define_flags() function centralizes definitions.

Testing

References: official/utils/testing/scripts

The …/scripts directory contains utilities for validating models and components during development and testing. It implements a standardized approach for running tests and checking models functionality across the TensorFlow models codebase.

The builds_common.sh file defines reusable Bash functions for common tasks like error handling, getting hardware info, filtering files, and diffing code changes. The ci_sanity.sh script serves as the driver for continuous integration tests. It runs linting checks by parsing the linting tool's output, counting errors, and returning a non-zero exit code if any non-whitelisted errors are found.

The presubmit.sh script defines validation checks that can be run locally before submitting code changes. It provides the py_test() function, which loops through Python test files like those under /official/, running each one sequentially and collecting exit codes to determine the overall test result. The py3_test() function calls py_test() specifically with the Python 3 interpreter. Together these functions implement a standardized way to validate code changes through testing.

Some other relevant subdirectories in /utils/testing include:

  • /testdata which contains sample inputs and outputs for tests
  • /metrics which defines implementations of common evaluation metrics that can be reused across models

Documentation

References: models

The …/docs directory contains utilities for generating documentation for the TensorFlow Models project. This includes functionality for generating API reference documentation on tensorflow.org.

The core functionality is provided by scripts. This script handles building documentation pages by iterating through source files and extracting comments marked with special tags. It handles parsing these comments into Markdown format and writing the output files.

Key aspects of the implementation include:

  • The script relies on functions to programmatically build documentation pages from source code.

  • Source files are discovered by walking the directories structure starting from the root path specified in a command line argument.

  • File parsing is done line-by-line using methods, which handles stripping comments and extracting fields marked with tags.

  • Extracted fields are converted to Markdown format strings by calling functions on the comment text.

  • Output files are written to the destination directory structure mirroring the source layout using functions. This allows linking between files on the generated site.

  • Configuration options like which tags to extract and the output directory are specified via flags parsed by functions from libraries.

  • Unit tests in files validate the core functionality of file parsing, field extraction, and Markdown formatting.

Configuration

References: models

Utilities for configuring models are primarily implemented through the Config class defined in …/__init__.py. This class provides a standardized way to define and access hyperparameters and other configuration settings for machine learning models.

The Config class supports initializing configurations from Python dictionaries, JSON or YAML files. It allows configurations to be overridden by specifying new values. Configurations can also be validated against constraints to check for invalid settings.

The Config class stores hyperparameters and other parameters as object attributes. Configuration objects defined as Config can then be passed to model building and training functions.

The ParamsDict class in …/params_dict.py acts as a container for hyperparameters organized in a nested dictionary structure. It provides functionality for initializing from files, accessing parameters, overriding values, and validating against constraints similar to the Config class.

Unit tests are provided in files like …/params_dict_test.py to validate the configuration functionality works as intended across different scenarios.

Distributed Training

References: models

The core functionality provided by the TensorFlow Models library for distributed training is contained in the …/common and …/core directories. These directories define utilities that standardize distributed training across models implemented in TensorFlow Models.

Some of the key functionality includes:

The …/distribute_utils.py file contains important functions for configuring distribution strategies for distributed training. The _mirrored_cross_device_ops() function in this file helps configure collective operations like all-reduces for MirroredStrategy by mapping the algorithm to classes like NcclAllReduce that perform the collective communication efficiently across multiple devices.

The get_distribution_strategy() function handles mapping hyperparameters to the correct configuration values for each distribution strategy, enabling flexible and consistent distributed training. This allows models to support all common distribution strategies with minimal code changes.

The core Trainer class defined in …/base_trainer.py inherits from _AsyncTrainer and handles both synchronous and asynchronous distributed training. It initializes properties like the model, optimizer, and metrics. Functions like train_step() run the task.train_step() method to perform distributed training across replicas, while updating metrics.

The InputReader class in …/input_reader.py provides a configurable pipeline for distributed input processing. It supports reading from files or TensorFlow Datasets via _read_data_source().

Callbacks

References: models

Keras callbacks allow customizing the training process. The …/keras_utils.py file contains utilities for working with Keras callbacks during training. Some important callbacks include:

  • The TimeHistory callback records training time and logs it to TensorBoard. Its on_epoch_end() method logs the time taken for each epoch.

  • The ModelCheckpoint callback saves model checkpoints periodically. It saves the model if the validation loss decreases, controlling overfitting.

These callbacks allow monitoring training progress and saving model checkpoints. They help improve model performance and reproducibility.

Synthetic Data

References: models

The …/model_helpers.py file contains utilities for generating synthetic data that can be used to test models. The synthetic_data() function generates random input-output pairs that mimic real data distributions. It supports different types of synthetic data generation through configuration options.

Some key functionality includes:

  • The synthetic_data() function takes hyperparameters like batch_size, input_shape, num_classes, label_smoothing to control data generation.

  • It generates random input tensors of shape input_shape filled with uniform noise using tf.random.uniform().

  • One-hot labels of shape num_classes are created from a categorical distribution parameterized by a probability vector. This vector is adjusted based on label_smoothing to add noise for regularization.

  • The inputs and labels are returned in a dictionary that can be directly used to train models, providing a standardized interface.

  • Additional options like dtype allow generating different data types like float16 for mixed precision testing.

  • The synthetic_data() function provides a simple but useful way to quickly generate random synthetic samples for tasks like unit testing models or debugging training loops without needing real datasets.

Datasets

References: research/vid2depth/dataset

The …/dataset directory contains functionality for loading, preprocessing, and converting datasets used by video-to-depth models. The core Dataset class defined in /__init__.py provides a standardized interface for handling different datasets.

The Dataset class handles loading raw data and preprocessing it based on configurable options. It loads data using dataset loaders defined in /dataset_loader.py, which abstract accessing different raw data sources. The Dataset then normalizes and batches samples before returning them.

Key files for datasets include:

  • /dataset_loader.py contains classes like KittiRaw() and Cityscapes() that load images and metadata from raw datasets.

  • /gen_data.py contains code to generate preprocessed training and validation data from raw datasets, saving examples efficiently to disk in a standardized format using the dataset loaders.

  • /__init__.py defines the core Dataset class.

  • /kitti contains KITTI data files and configuration, including lists of test images/scenes in test_files_eigen.txt and test_scenes_eigen.txt.

The Dataset class and dataset loaders allow generating preprocessed data in a source-agnostic way. This standardized interface and preprocessing functionality supports loading different datasets for video-to-depth models.

Research Datasets

References: models

The core functionality for loading and preprocessing research datasets is contained in the …/slim directory. This directory provides standardized implementations of loading common computer vision and natural language processing datasets into TensorFlow records and example protos that can then be used across different models implemented with TensorFlow Slim.

The Dataset class defined in …/__init__.py acts as the primary interface for defining datasets. Subclasses of Dataset must implement methods like _parse_function() that specify how to parse individual examples (e.g. decoding and preprocessing images) respectively.

Some important datasets implemented include ImageNet, CIFAR-10, MNIST, and popular NLP datasets. The data loading and preprocessing pipelines are encapsulated within subclasses located in subdirectories. This allows new datasets to be easily added by defining a new Dataset subclass without having to modify model code.

The _parse_function() method on Dataset subclasses handles key preprocessing steps. For images, this typically involves functions like to randomly augment images on-the-fly in a way that is consistent across epochs. Text examples are parsed.

Official Datasets

References: models

The core functionality for loading and preprocessing datasets for official TensorFlow models is contained in the …/common directory. This directory provides utilities for loading common computer vision and natural language processing datasets and converting them into standardized TensorFlow formats that can be easily used for model training and evaluation.

The main interface for defining datasets is the Dataset class. The Dataset class provides a consistent way to represent different datasets so they can be loaded and preprocessed in a uniform manner. The Dataset class is initialized with configuration options and handles important functionality like shuffling examples, batching examples into minibatches of a fixed size, and iterating through examples in an efficient manner that is compatible with TensorFlow.

An important utility function for selecting the appropriate dataset loading function based on file type is the pick_dataset_fn function. This function avoids duplicating dataset loading code for different file types by selecting either a partial function based on whether files are in a plain format or compressed format like gzip. Using partial functions allows pick_dataset_fn to return callables that can be used to instantiate TensorFlow datasets, providing flexibility.

The Dataset class relies on preprocessing functions defined in subdirectories like …/dataloaders to handle important steps like decoding, parsing and normalizing input data into standardized tensors expected by the models. For example, preprocessing steps for image classification are handled by the Parser class. The Dataset class also uses the InputReader class to handle input processing and distributing data across devices in a flexible way.

Video-to-Depth Datasets

References: models

The …/dataset directory contains functionality for loading and preprocessing popular datasets for video-to-depth research. This includes the KITTI and Cityscapes datasets, which provide stereo video sequences paired with ground truth depth maps.

The Dataset class provides a common interface for defining video-to-depth datasets. Subclasses of Dataset implement loading of specific datasets.

Business logic:

  • The Dataset class provides a standardized way of accessing samples from different datasets. Subclasses load samples by overriding methods to return preprocessed sample tuples containing images and depths.

Training

References: official/core, research/slim/deployment

Training components are handled through the Trainer class defined in …/base_trainer.py. The Trainer initializes important properties like the model, optimizer, metrics, and checkpoint. It handles distributing datasets and running the core training and evaluation loops through functions like train_step(), eval_step(), train_loop_end(), and eval_end().

The train_step() function runs the task.train_step() method to perform the training step. This updates the model weights and calculates metrics like loss. eval_step() similarly runs task.validation_step() to evaluate the model. Functions like train_loop_end() and eval_end() consolidate metrics, reset them, and return results. The Trainer handles common training and evaluation logic, allowing tasks to focus on defining their specific models, losses, and metrics.

Optimization is handled through the optimizer initialized in the Trainer. The optimizer is configured based on the experiment configuration using configure_optimizer() defined in base_trainer.py. This function sets properties like loss scaling for mixed precision if specified in the config. During training, the optimizer updates the model weights based on the gradients calculated in train_step().

Distributed training is supported using any TensorFlow distribution strategy. The Trainer initializes the strategy and distributes datasets using distribute_dataset(). It builds the model, loss, and metrics to be compatible with the strategy using strategy.run(). Functions like train_step() and eval_step() then execute the distributed operations such as applying gradients and running the train op across all devices in the cluster. This allows training large models on multiple GPUs or other accelerators.

The OrbitExperimentRunner defined in train_lib.py orchestrates running complete experiments. It initializes important components like the Trainer, Controller, and CheckpointManager. The Controller's methods like train() and evaluate() then execute the full training and evaluation loops by running distributed operations across all devices and hosts configured by the strategy.

Modeling

References: research/slim/nets

This section covers building blocks for constructing models like layers, cells, backbones, heads, and optimizers. These components provide reusable pieces that can be assembled to build complex neural network architectures.

The …/conv_blocks.py file defines common convolutional building blocks for constructing efficient convolutional models like MobileNet. The key block is expanded_conv() which implements the "inverted bottleneck" block of MobileNetV2. It supports options like splitting, residuals and squeeze-excite attention. These modular blocks can then be composed to define model architectures.

The …/resnet_utils.py file contains utilities for building ResNet models. It defines the Block class, which represents a single block or unit within network architectures. For ResNet it defines a residual block, encapsulating its scope, unit function and arguments. This provides a clean way to represent repeating blocks that make up models. The file also contains functions for operations like subsampling and strided convolutions used when stacking blocks.

The …/mobilenet_v1.py file contains an object-oriented implementation of the core MobileNet architecture in the mobilenet_base() function. It iterates through convolution definitions to construct the network topology, applying operations to control aspects like output stride. Utilities like global_pool() and training_scope() set up the classification head and training functionality. The file also defines common convolutional layers used in MobileNet like Conv and DepthSepConv named tuples.

The …/mobilenet_v2.py file defines the MobileNetV2 model using the expanded_conv block. It implements the published MobileNetV2 specification by composing these blocks and applying techniques like depth multiplication, squeeze-excite, and fused convolutions according to the specification.

The …/mobilenet_v3.py file defines implementations of MobileNetV3 models like MobilenetV3Large using operations like mbv3_op that incorporate techniques like squeeze-excite and fused convolutions. It constructs these models by composing these operations according to the published specifications.

The …/inception_utils.py file contains utilities for building inception models. It defines the inception_block_v1_3d() function, which extends the standard Inception block to 3D. This allows modeling spatial-temporal features. The file also contains functions for operations like 3D convolutions and normalization used across inception models.

Evaluation

References: research/deeplab/evaluation

This section covers metrics, analysis, and visualization tools for rigorously evaluating machine learning models. Key functionality includes implementations of important metrics, utilities for offline and online evaluation, and visualization of results.

The core metrics are defined in classes inheriting from SegmentationMetric in …/base_metric.py. This includes PanopticQuality in …/panoptic_quality.py for measuring panoptic segmentation accuracy, and ParsingCovering in …/parsing_covering.py for semantic segmentation covering.

PanopticQuality computes the Panoptic Quality metric by comparing predicted and ground truth category and instance labels using _ids_to_counts to count label intersections. It updates per-category statistics in compare_and_accumulate and calculates the final PQ score from these in result_per_category and result.

ParsingCovering measures semantic segmentation covering by calculating maximum IoUs between predictions and ground truths in compare_and_accumulate using NumPy operations. It accumulates weighted IoU and area statistics used by result_per_category and result to compute the final PC score.

Offline evaluation is supported by eval_coco_format in …/eval_coco_format.py, which takes predictions and annotations to run metrics like PQ and PC. Online "streaming" metrics are implemented in streaming_metrics using TensorFlow variables to aggregate results across a dataset stream.

Sample data for testing is provided in Test Data. Thorough unit tests validate the metrics in files like panoptic_quality_test.py and streaming_metrics_test.py. Visualization and debugging tools are implemented but not described in the provided summaries.

Training

References: orbit, official/core, research/slim/deployment

The section 'Training' covers the core components used for implementing training loops, optimization, and distributed training functionality. Key aspects covered include:

The Trainer class defined in …/base_trainer.py handles common training loop functionality. It initializes important properties like the model, optimizer, and metrics. Functions like train_step() run the task.train_step() method to perform the optimization step, while updating metrics. eval_step() similarly runs evaluation. train_loop_end() and eval_end() consolidate metrics and return results. The Trainer handles common training and evaluation logic so tasks can focus just on their models, losses, and metrics.

The StandardTrainer class provides a standardized implementation of a training loop on top of the base AbstractTrainer class. It structures a training loop by calling train_loop_begin(), then looping over calls to the abstract train_step() method defined by subclasses, and finally calling train_loop_end().

The SingleTaskTrainer class handles training a single-output model. The core of the training is implemented in train_step(), which uses a GradientTape to calculate gradients, calls the loss function, applies gradients with the optimizer, and updates metrics like self.train_loss.

The StandardEvaluator class similarly provides a standard implementation of an evaluation loop on top of AbstractEvaluator.

The SingleTaskEvaluator class extracts targets, runs the model via its eval_step() method, and updates TensorFlow Keras metrics which are returned via eval_end().

The OrbitExperimentRunner class in …/train_lib.py uses the TFM orbit library to define the overall training loop for running experiments.

Optimization

References: orbit, official/core/base_trainer.py

The core components for implementing optimization and training loops in TensorFlow models are provided by the StandardTrainer class defined in …/standard_runner.py.

The StandardTrainer class structures the training loop by calling train_loop_begin(), then looping over calls to the abstract train_step() method, and finally calling train_loop_end(). This provides a standardized way for models to implement their specific single step optimization logic via the train_step() method while handling the overall loop structure.

The train_step() method is where models define their single optimization step. This typically involves using a GradientTape to calculate gradients, applying the gradients to variables with an optimizer like Adam or SGD, and updating metrics with the loss value. By encapsulating the single step logic here, models can focus on just their optimization calculations while leveraging the full training loop provided by StandardTrainer.

Some key aspects of how StandardTrainer implements the loop include:

  • It supports efficiently running the loop using either TensorFlow functions for high performance, or a basic Python loop
  • Metrics are reset before each loop via train_loop_begin() and the final values exported after via train_loop_end()

This provides models a standardized, high-performance way to define their optimization logic while leveraging the full training loop functionality and optimizations provided by StandardTrainer.

Distributed Training

References: research/slim/deployment/model_deploy.py

The …/model_deploy.py file provides utilities for distributed training of models across multiple GPUs and machines. It implements common algorithms for distributed training such as:

  • Model replication: The create_clones() function replicates the model function multiple times, assigning each clone to a separate device. This allows data parallelism by computing gradients on multiple devices simultaneously.

  • Gradient aggregation: In the optimize_clones() function, the losses and gradients are computed for each clone separately. The gradients are then summed across clones before being applied to the optimizer. This aggregates the distributed gradients.

  • Configuration: The DeploymentConfig class controls placement of variables, clones, queues, and other objects for distributed training. It specifies properties like the number of clones and devices each clone is assigned to.

  • Deployment: The deploy() function handles full deployment of the model using the configuration. It runs the training loop, collecting summaries and checkpointing the model.

The key classes and functions implemented are:

  • DeploymentConfig: Configures properties like number of clones and their assigned devices.

  • create_clones(): Replicates the model function and assigns clones to different devices using the configuration.

  • optimize_clones(): Computes losses and gradients for each clone separately, then sums the gradients across clones before applying to the optimizer.

  • deploy(): Deploys the full training loop across devices using the configuration.

Callbacks

References: official/utils/misc/keras_utils.py

The Keras callbacks defined in …/keras_utils.py allow monitoring and logging metrics during training. The TimeHistory callback calculates steps/second, examples/second, and epoch runtimes, logging these periodically or saving summaries to TensorBoard. It tracks timing data by overriding lifecycle methods like on_batch_end() and on_epoch_begin(). The SimpleCheckpoint callback saves checkpoints by calling save() on a tf.train.CheckpointManager passed to its constructor. It saves after each epoch via on_epoch_end(). These callbacks provide a standardized way to monitor training progress and save model checkpoints across different models and experiments.

The TimeHistory callback calculates timing metrics if enough steps have elapsed. It logs the metrics or saves TensorBoard summaries depending on if a SummaryWriter is provided. TimeHistory tracks start and end times by overriding methods like on_epoch_begin() and on_train_end().

The SimpleCheckpoint callback saves checkpoints by calling the save() method on a CheckpointManager object passed to its constructor. It saves models after each epoch completes by implementing checkpoint saving logic in the on_epoch_end() method.

Checkpointing

References: official/core/actions.py

The core functionality for checkpointing and saving/restoring models is handled by several classes in …/actions.py.

The EMACheckpointing class handles saving exponential moving average (EMA) checkpoints during evaluation. It requires an ExponentialMovingAverage optimizer from TensorFlow and uses the optimizer's swap_weights() method to save the averaged weights to disk. The averaged weights are then swapped back after saving. This allows checkpointing models that use EMA for evaluation.

The RecoveryAction class implements recovery from checkpoints if the loss blows up during training. It simply loads the previous checkpoint on a loss threshold being exceeded, as checked by the RecoveryCondition class. The RecoveryCondition checks the loss value against a configurable threshold.

These action classes are configured and returned by the get_eval_actions() and get_train_actions() functions based on the experiment configuration and trainer properties. The functions also add recovery actions if enabled.

Metrics

References: official/core/base_task.py

The Task class in …/base_task.py provides the core functionality for tracking metrics during training. The Task implements abstract methods that tasks must define to build their metrics. This includes the build_metrics() method, which tasks use to define the metrics that will be computed during training and validation.

The Task class handles computing the metrics during the train_step() and validation_step() methods. These methods call the task's metric functions and aggregate the results. Metric values are returned and can be logged or used for early stopping. The Task also provides utility methods for distributing the computation across devices and hosts during distributed training.

Two important components for metrics are the Metric class and the MetricFn class. The Metric class represents a single metric like accuracy or loss. It accumulates the metric value during steps and computes the final result. The MetricFn class wraps a TensorFlow metric function, handles resetting the metric at the start of each step, and computing the final metric value. Tasks return Metric objects from their build_metrics() method to integrate the metrics into the training/validation steps.

Visualization

References: models

The code provides tools for visualizing models and results during training. Some key functionality includes:

Visualizations help debug and improve models by allowing inspection of predictions overlaid on inputs. Color coded outputs clearly show predictions compared to ground truths on a per-pixel basis.

Modeling

References: official/modeling, research/slim/nets

The …/modeling directory contains implementations of common components for constructing neural network models in TensorFlow. It provides modular building blocks that can be composed to build complex state-of-the-art models across domains.

Some key aspects included are neural network layers, cells that combine layers, backbone networks, heads or task-specific components, and utilities for tasks like serialization and distribution. These building blocks handle common needs when defining models so they can focus on their task-specific logic.

The …/activations subdirectory contains implementations of activation functions like ReLU, GELU, and Mish that are used across many model types. These are implemented as differentiable TensorFlow operations.

The …/optimization subdirectory provides classes for optimization algorithms. For example, the OptimizerFactory class handles constructing optimizers like SGD and Adam from configuration objects. The LAMB class implements the layer-adaptive optimizer.

The …/hyperparams subdirectory defines reusable classes for representing hyperparameters. The Config and ParamsDict classes standardize how hyperparameters are accessed and validated throughout model building.

Now we will discuss some important classes and functions in more detail:

The Block class defined in files like …/resnet_utils.py represents a single repeating block or unit within network architectures. For example, in ResNet it defines a residual block, encapsulating its scope and residual unit function. This provides a clean way to represent repeating blocks that make up models.

The resnet_v2() function in …/resnet_v2.py demonstrates how Block objects can be composed to build full models. It takes a list of Block objects as input and stacks their units using resnet_utils.stack_blocks_dense() while controlling aspects like output stride. This modular approach makes architectures easy to define and extend.

The bottleneck() function implements the basic residual block unit used throughout ResNet models. It applies the preactivation structure - batch norm, ReLU, convolutions, addition of input - that is critical to how ResNet implements residual learning.

Layers

References: models

The …/modeling directory contains implementations of common neural network layers that can be composed together to build complex models.

The BaseLayer class provides a base class for building quantized model layers. Its __call__() method applies layers sequentially. It handles common functionality like adding weights and biases, regularization losses, and methods for weight quantization like quantize_parameter().

The Conv2D layer applies a 2D convolution to the input, supporting specifying the kernel size, strides, padding, and other arguments.

The Embedding layer looks up embeddings for each token id from a trainable weights table. It supports specifying the vocabulary size, embedding dimension, and other properties.

The BatchNormalization layer performs batch normalization by centering and scaling inputs across a batch dimension.

The AveragePooling2D layer takes the average value in each pool region, supporting specifying the pool size and strides.

The LSTM layer implements the LSTM recurrent cell type, supporting specifying the units and other arguments.

Optimizers

References: official/modeling/optimization

The …/optimization package contains implementations of common optimization algorithms used for training neural networks. The key classes include LAMB which implements popular gradient-based optimization methods.

The LAMB optimizer in …/lamb.py inherits from Optimizer and implements the layer-adaptive moment estimation algorithm. It maintains moment estimate slots for variables and computes a scaled learning rate ratio based on weight and gradient norms for layer adaptation.

The OptimizerFactory in …/optimizer_factory.py handles constructing optimizers from configuration objects. It maps configuration types like "sgd" to the appropriate class implementations.

Thorough unit tests for the optimizers are provided in files like …/lamb_test.py and …/optimizer_factory_test.py to validate their functionality.

Activations

References: official/modeling/activations

The …/activations directory contains implementations of common activation functions used in neural networks. Activations like ReLU, GELU, and Mish introduce non-linearities into neural network models which help them learn complex patterns from data.

Each activation is implemented as a simple and efficient TensorFlow operation. ReLU clamps values to be non-negative using tf.nn.relu6, while GELU leverages the existing Keras implementation. Mish directly calculates the softplus and tanh components to output values between -1 and 1. These activations can be used interchangeably as network building blocks.

The activations are tested to ensure they match reference implementations. For example, ReLU is compared to tf.nn.relu6, while Mish and GELU are validated on sample inputs. This confirms the activations are correctly implemented and differentiable.

The …/__init__.py file collects the common activations in one place. It imports functions like ReLU, GELU, Mish, hard_sigmoid and hard_swish before re-exporting them. This provides a clean interface for importing any activation without needing to know its specific location.

Hyperparameters

References: official/modeling/hyperparams

The ParamsDict class provides a standardized way to configure hyperparameters during experiments. ParamsDict allows hyperparameters to be organized in a nested dictionary-like structure and supports operations like overriding values, locking the structure to prevent changes, and validating consistency between parameters based on defined restrictions.

The core class defined for handling hyperparameters is ParamsDict. It initializes from a default dictionary and allows parameters to be accessed and modified via attribute syntax like a dictionary. Methods like override(), lock(), and validate() allow changing values, preventing further changes, and checking restrictions.

The override() method recursively overrides nested parameters while respecting whether new keys should be allowed. It uses the _override() private method to contain the core recursive logic.

The validate() method checks user-defined restrictions specified as strings against the parameter values. It raises errors if any inconsistencies are found. validate() uses the _get_kv() and _get_kvs() helper functions to extract key-value pairs from the restriction strings.

Utility functions include read_yaml_to_params_dict() to initialize hyperparameters from YAML files, save_params_dict_to_yaml() to write hyperparameters back out to YAML, and override_params_dict() which tries various methods to override hyperparameters from different sources like dictionaries, strings or files.

The nested_csv_str_to_json_str() function is important for parsing hyperparameters configured as comma-separated strings. It converts these strings to JSON strings, handling nesting via dot notation and supporting quoted strings. This function uses the _PARAM_RE regular expression to extract the key-value pairs.

Utilities

References: official/modeling/tf_utils.py, official/modeling/__init__.py

The …/__init__.py file collects common utilities that are used throughout the modeling module to build neural networks. This includes common layers like Dense and Dropout that are used as building blocks, the Model base class for defining models, and preprocessing utilities in image/image_preprocessing.

The Model class standardizes the API for model training and evaluation. Models inherit from this class and must implement the call() method for the forward pass. The preprocessing utilities help apply common transformations to input data like images before feeding to models.

Some key utilities provided for modeling include:

  • Serialization functionality allows saving models to config dictionaries using serialize_keras_object() and loading from configs with deserialize_keras_object(). This supports persistently storing model architectures.

  • Initialization utilities like clone_initializer() ensure each layer gets its own copy of initialization functions like VarianceScaling instead of sharing references, for deterministic initialization.

  • Activation functions are retrieved in a standardized way via get_activation() from …/tf_utils.py. This maps names to actual TensorFlow/Keras functions and also supports using Keras layers as activations.

  • Shape and dimension utilities like get_shape_list() and assert_rank() help work with tensors and ensure they have expected properties.

Multi-Task Learning

References: official/modeling/multitask

The core functionality for multi-task learning with shared models is encapsulated in several key classes defined in the …/multitask directory. The MultiTaskBaseModel class serves as an abstraction for multi-task models, allowing them to share portions of the model architecture while also defining task-specific heads.

The MultiTaskBaseModel holds the computation for different sub-tasks by initializing a dictionary of sub-task models in its _instantiate_sub_tasks() method. This method must be implemented by subclasses to define the specific sub-tasks and return a dictionary mapping task names to tf.keras.Model objects. The sub_tasks property then provides access to this dictionary, exposing the individual task models.

The MultiTask class manages the set of tasks by initializing with a dictionary of Task objects. It handles running joint training steps across all tasks simultaneously via its joint_train_step() method. This method forwards inputs to the appropriate sub-tasks, calculates losses for each task, applies weights, and sums for the total loss in one training step.

The MultiTaskBaseTrainer handles training a MultiTaskBaseModel using a specified MultiTask and Optimizer. It overrides methods like train_step() to call the MultiTask's joint_train_step() and implement the shared training logic. This trains all tasks jointly while sharing portions of the model architecture defined in the base model.

The MultiTaskEvaluator allows evaluating multiple tasks on a multi-task model simultaneously. It builds distributed datasets for each task, constructs validation metrics, and runs evaluation in a coordinated way, returning results for each task.

Privacy

References: official/modeling/privacy

The DifferentialPrivacy class in …/__init__.py implements the core functionality for differentially private training of TensorFlow models. The class allows configuring the privacy mechanism and noise levels to be applied during training.

The DifferentialPrivacy class takes a DifferentialPrivacyConfig as an argument to its constructor. This config, defined in …/configs.py, encapsulates the key hyperparameters for differential privacy training - the clipping norm and noise multiplier. Using a config class provides a clean interface for configuring privacy without passing many arguments around.

During the training process, the DifferentialPrivacy class applies gradient clipping and noise addition to model updates before applying them. It uses the operations defined in …/ops.py - clip_l2_norm() clips gradients to the global norm specified in the config, while add_noise() adds mean-zero normal noise scaled by the noise multiplier to each gradient. These core operations implement the differential privacy mechanism.

Unit tests in …/configs_test.py and …/ops_test.py extensively test that the config class and privacy operations behave as expected. For example, tests check default config values, that clipping limits gradients correctly, and noise is added deterministically for testing purposes. Thorough testing helps prevent regressions in the differential privacy functionality.

Fast Training

References: official/modeling/fast_training

Fast training techniques implemented in the code include progressive training and horizontal model expansion. Progressive training gradually increases model complexity through multiple training stages. The ProgressivePolicy class defines the core interface that policies must implement for progressive training. It provides methods like get_model() to return the model for each stage. The ProgressiveTrainer handles running the training loop over each stage by calling the policy interfaces. Checkpointing is supported between stages using CheckpointWithHooks, allowing objects like the model to be seamlessly replaced.

The …/progressive directory contains the main implementation of progressive training. It defines the ProgressivePolicy abstract base class that subclasses must implement to define the specifics of each stage. The ProgressiveTrainer class implements the overall training loop, initializing using the policy and training over each stage. Tests in …/trainer_test.py validate functionality works correctly between stages.

Horizontal model expansion allows training a wider model initialized from a pretrained narrow model. The …/experimental directory contains utilities for expanding models. Functions like expand_vector() and expand_1_axis() duplicate weights while adding noise to prevent degenerate solutions. model_to_model_2x_wide() iterates through a narrow model, expands each variable using var_to_var(), and assigns to a wider target model. This preserves the input-output relationship. Tests in …/tf2_utils_2x_wide_test.py validate the expansion functions.

Evaluation

References: research/deeplab/evaluation

This section covers tools for analyzing and visualizing models during and after training. Key functionality includes metrics, visualization, and debugging.

The …/evaluation directory contains important metrics for semantic and panoptic segmentation. The SegmentationMetric base class in base_metric.py defines the interface for segmentation metrics. Subclasses like PanopticQuality and ParsingCovering implement important metrics like Panoptic Quality (PQ) and Parsing Covering (PC) by inheriting from SegmentationMetric.

PanopticQuality computes PQ by comparing predicted and ground truth category and instance labels in its compare_and_accumulate method. It stores per-category statistics and computes the final PQ score from these in result_per_category and result. ParsingCovering similarly implements PC, calculating maximum IoUs and accumulating weighted IoU and area statistics to derive the final PC score.

The streaming_metrics.py module contains functions for evaluating metrics in an online streaming fashion over multiple examples using TensorFlow. _running_total() maintains running sums between calls using a TensorFlow variable. streaming_panoptic_quality() and streaming_parsing_covering() compute PQ and PC respectively by accumulating results across examples.

Thorough unit tests for metrics and evaluation components are provided in files such as …/panoptic_quality_test.py and …/streaming_metrics_test.py. These validate the metrics are computed correctly.

Metrics

References: research/deeplab/evaluation/base_metric.py, research/deeplab/evaluation/panoptic_quality.py, research/deeplab/evaluation/parsing_covering.py

The core metrics implemented in this code are SegmentationMetric, PanopticQuality, and ParsingCovering. SegmentationMetric defines the base interface for segmentation metrics, providing common functionality like comparing predictions to ground truth and accumulating results.

PanopticQuality inherits from SegmentationMetric and implements the Panoptic Quality (PQ) metric. PQ measures how well a panoptic segmentation matches the ground truth in both categories and instances. Its compare_and_accumulate method computes intersections between predictions and ground truth using _ids_to_counts, calculates IoU for true matches, and updates per-category counts. From these counts, result_per_category calculates per-category PQ and result averages this to get the overall PQ score. detailed_results breaks results down into things/stuff subsets.

ParsingCovering also inherits from SegmentationMetric and implements the Parsing Covering metric. It computes maximum IoU for each ground truth instance using NumPy operations on label arrays. compare_and_accumulate accumulates weighted IoU and area statistics, while result_per_category returns the weighted IoU for each category. result averages per-category results to compute the final Parsing Covering score.

The key classes implement important segmentation metrics. SegmentationMetric defines the interface, while PanopticQuality and ParsingCovering contain the core logic for the PQ and Parsing Covering algorithms through methods like compare_and_accumulate and result.

Utilities

References: research/deeplab/evaluation/eval_coco_format.py, research/deeplab/evaluation/streaming_metrics.py

This section covers utilities for performing model evaluation both offline and online. The key utilities provided are for offline evaluation, where metrics are computed by processing all examples in a dataset at once, and online or streaming evaluation where metrics are accumulated incrementally over examples.

The …/eval_coco_format.py file contains the main logic for offline evaluation. It uses the _matched_annotations function to iterate over annotation pairs and extract category and instance maps using _split_panoptic. Metrics like PQ, SQ and RQ are then computed by passing these maps to _panoptic_quality_helper. This file also supports parallelizing the evaluation across multiple processes.

Streaming evaluation is supported by utilities in …/streaming_metrics.py. This allows accumulating metrics over many examples without needing to process all data at once. The _running_total() function maintains running totals of metric values using TensorFlow variables. Metrics like PQ, SQ and RQ can then be computed in a streaming fashion by the streaming_panoptic_quality() function, which calls _panoptic_quality_helper() on each example and accumulates the results. Similarly, the streaming_parsing_covering() function computes metrics like weighted IoU incrementally. By accumulating metrics this way, very large datasets can be evaluated without requiring loading all data into memory at once.

Test Data

References: research/deeplab/evaluation/testdata

The …/testdata directory contains sample input images that can be used for testing semantic image segmentation models. It holds documentation in the file …/README.md which describes the test images included.

The README file documents that the test data directory contains 4 sample RGB images - team_input.png, cat_input.jpg, bird_input.jpg, and congress_input.jpg. These images cover different subjects like people, animals and scenery. The sources of the images are cited in the README for reference.

No code implementation or algorithms are directly included in this test data directory. It solely contains sample input content that can be used to validate that a semantic segmentation model is functioning as intended on new images. Developers can load these images and check that the model produces expected output segmentation masks. The Evaluation section covers how model predictions are scored against ground truth data.

Unit Tests

References: research/deeplab/evaluation/panoptic_quality_test.py, research/deeplab/evaluation/streaming_metrics_test.py

The unit tests for the evaluation components are located in two main files: …/panoptic_quality_test.py and …/streaming_metrics_test.py. These files contain rigorous tests that validate the core evaluation metrics are calculating as expected.

The panoptic_quality_test.py file tests the PanopticQuality class, which calculates the panoptic quality score by comparing a predicted segmentation to the ground truth. It accumulates statistics like true positives, false positives, false negatives for each category. Various test cases are created to validate the class handles different matching scenarios between predictions and ground truth correctly. Real prediction and ground truth data is also loaded and tested.

The streaming_metrics_test.py file contains two test classes that evaluate the streaming metrics for panoptic quality and parsing covering. These metrics calculate important scores like PQ, SQ, RQ and covering/IoU in an online fashion on streams of images. The test classes initialize local variables, run the metric update ops on test data in TensorFlow sessions, and assert the final results match expected values.

Both files thoroughly validate the core evaluation functionality through unit tests. A variety of test inputs exercise different code paths to ensure the metrics are calculated accurately. Placeholders and real data help test the end-to-end pipelines. This confirms the metrics can be reliably used to evaluate panoptic segmentation models.

Formats

References: research/deeplab/evaluation/test_utils.py

This section details utilities for loading common test data formats used in evaluation. The …/test_utils.py file contains functions that make it convenient to read data formats used for testing evaluation code.

The read_test_image() function loads images from the …/testdata directory as NumPy arrays using scipy.misc.imread(). This provides a simple way to read images for tests.

The read_segmentation_with_rgb_color_map() function reads semantic segmentations stored as labeled RGB images. It first loads the RGB image with read_test_image() then assigns semantic labels to each pixel by looking up its color in the provided rgb_to_semantic_label mapping. This validates all colors are present in the mapping to ensure loaded data matches expectations.

The panoptic_segmentation_with_class_map() function loads panoptic segmentation data stored with an instance segmentation map and instance_label_to_semantic_label mapping. It loads the instance map using read_test_image() in grayscale mode then looks up the class label for each pixel's instance ID in the mapping. This loads and validates common panoptic segmentation test formats.

Overall, these utilities make it simple to load test data stored in RGB semantic, instance map, and panoptic segmentation formats while validating the data matches provided mappings. This allows convenient writing of unit tests for evaluation code on real data samples.