TensorFlow Models — Administrator & Developer Guide

1. Architecture Overview

TensorFlow Models is Google's official repository of machine learning model implementations for TensorFlow 2.x. The codebase is organized into three main pillars: official (production-quality, Google-maintained models), research (paper implementations, community-maintained), and orbit (a standalone training loop library). Together, they form a layered framework where official/core provides the configuration, registry, and training infrastructure; domain-specific libraries (official/vision, official/nlp, official/recommendation) provide task implementations and model architectures; and Orbit manages the outer training loop with distribution strategy support.

The pip-installable package tf-models-official (v2.20.0) bundles the official/ and orbit/ directories. Research models are used from source only. All models are built on tf_keras (the standalone Keras fork) and TensorFlow 2.x, with native support for tf.distribute strategies including MirroredStrategy, TPUStrategy, MultiWorkerMirroredStrategy, and ParameterServerStrategy.

Architecture Diagram

┌─────────────────────────────────────────────────────────────────────┐
│                     tensorflow_models (pip API)                     │
│   tfm.vision  │  tfm.nlp  │  tfm.uplift  │  tfm.core/optimization │
├───────────────┴───────────┴──────────────┴─────────────────────────┤
│                      official/ (core library)                       │
│  ┌──────────┐  ┌──────────┐  ┌───────────────┐  ┌──────────────┐  │
│  │  core/   │  │ modeling/ │  │    common/    │  │  utils/flags │  │
│  │ registry │  │ hyperpar │  │ distribute_   │  │  _base.py    │  │
│  │ task_fac │  │ optimiz. │  │ utils.py      │  │  _perf.py    │  │
│  │ train_lib│  │ grad_util│  │ dataset_fn.py │  │              │  │
│  └────┬─────┘  └────┬─────┘  └───────┬───────┘  └──────────────┘  │
│       │              │                │                             │
│  ┌────▼──────────────▼────────────────▼──────────────────────────┐  │
│  │                  Domain Libraries                             │  │
│  │  vision/        nlp/          recommendation/    projects/    │  │
│  │  (ResNet,       (BERT,        (ranking, uplift)  (YOLO,      │  │
│  │   RetinaNet,    ALBERT,                          CenterNet,  │  │
│  │   MaskRCNN,     Longformer)                      DeepLab,    │  │
│  │   DeepLab)                                       SimCLR...)  │  │
│  └──────────────────────────────────────────────────────────────┘  │
├───────────────────────────────────────────────────────────────────┤
│  orbit/                              │  research/                 │
│  Controller (train/eval loops)       │  object_detection/         │
│  StandardTrainer/StandardEvaluator   │  deeplab/ slim/ delf/      │
│  Actions (EMA, Export, Pruning)      │  audioset/ adversarial_text│
│  SummaryManager, CheckpointManager   │  (standalone, from source) │
└──────────────────────────────────────┴───────────────────────────┘
                            │
                   ┌────────▼────────┐
                   │  tf.distribute  │
                   │  Strategy Layer │
                   │  (GPU/TPU/Multi)│
                   └─────────────────┘

Component Inventory

Technology Stack

Note

The codebase uses import tensorflow as tf, tf_keras throughout — this dual import pattern ensures compatibility with TensorFlow 2.16+ where Keras was separated into tf_keras. All Keras API calls use tf_keras rather than tf.keras.

2. System Requirements

Software Dependencies (from official/requirements.txt)

3. Configuration Reference

Configuration System Architecture

TensorFlow Models uses a hierarchical configuration system:

1. Python dataclasses (official/core/config_definitions.py) — Define ExperimentConfig, RuntimeConfig, TrainerConfig, and task-specific configs as typed, validatable Python objects
2. YAML files — Serialized experiment configs in configs/experiments/ directories
3. Gin-config — Additional binding overrides for flexible parameterization
4. CLI flags — Runtime overrides via --params_override and --config_file
5. ParamsDict — Legacy config system used by official/legacy/ models (official/modeling/hyperparams/params_dict.py)

Config Files

Environment Variables

CLI Flags — Common Training Flags

These flags are registered by official/utils/flags/_base.py and available in most training scripts:

CLI Flags — Vision Training (official/vision/train.py)

CLI Flags — Legacy Detection (official/legacy/detection/main.py)

Legacy Detection Base Configuration

The base config (official/legacy/detection/configs/base_config.py) defines defaults for all detection models:

Important

The frozen variable prefix RESNET_FROZEN_VAR_PREFIX freezes conv1 and conv2_x layers during fine-tuning. This speeds up training and slightly improves accuracy. Set train.frozen_variable_prefix to this regex when loading from pre-trained ResNet checkpoints.

Distribution Strategy Configuration

# From official/common/distribute_utils.py
distribution_strategy = distribute_utils.get_distribution_strategy(
    distribution_strategy='mirrored',   # Strategy type
    all_reduce_alg='nccl',              # All-reduce algorithm
    num_gpus=4,                         # GPUs per worker
    tpu_address='grpc://10.0.0.1:8470'  # TPU address (for TPU strategy)
)

Warning

When specifying num_gpus > 0 with distribution_strategy=off, a ValueError is raised. To use GPUs without a distribution strategy, set num_gpus=0 and manage device placement manually.

4. Deployment & Scaling

pip Package Distribution

The official/pip_package/setup.py builds two packages:

The nightly version bumps the minor version and appends a date suffix.

Building from Source

cd official/pip_package
python setup.py sdist bdist_wheel
pip install dist/tf_models_official-2.20.0-py3-none-any.whl

The package includes all official/ code except research/, official.pip_package, official.benchmark, official.colab, and official.recommendation.ranking.data.preprocessing. Test files (*_test.py) are excluded from distribution.

Multi-Worker Distributed Training

# Worker 0
export TF_CONFIG='{"cluster":{"worker":["host1:12345","host2:12345"]},"task":{"type":"worker","index":0}}'
python official/vision/train.py --experiment=resnet_imagenet --config_file=config.yaml

# Worker 1
export TF_CONFIG='{"cluster":{"worker":["host1:12345","host2:12345"]},"task":{"type":"worker","index":1}}'
python official/vision/train.py --experiment=resnet_imagenet --config_file=config.yaml

The configure_cluster() function in distribute_utils.py reads TF_CONFIG, extracts worker_hosts and task_index, and configures the appropriate strategy. When the number of workers is >= 2, multi-host mode activates and batch sizes are automatically divided across replicas.

TPU Training

python official/vision/train.py \
  --experiment=resnet_imagenet \
  --mode=train_and_eval \
  --model_dir=gs://bucket/model_dir \
  --config_file=config.yaml \
  --params_override="runtime.distribution_strategy=tpu,runtime.tpu=grpc://${TPU_IP}:8470"

When strategy_type == 'tpu', the system automatically:

• Enables use_bfloat16 for mixed precision
• Enables use_sync_bn for synchronized batch normalization
• Adjusts batch partitioning for TPU cores

Docker Deployment — Waste Identification API

The waste identification project provides a FastAPI-based prediction server:

# From official/projects/waste_identification_ml/docker_solution/prediction_api/
FROM tensorflow/tensorflow:latest
COPY . /app
WORKDIR /app
RUN pip install fastapi uvicorn Pillow
EXPOSE 5000
CMD ["python", "app.py"]

The server loads Mask R-CNN models at startup (model_manager.load_all_models()), accepts image uploads via POST /predict, resizes to 512×1024, and returns JSON detection results. Runs on Uvicorn (ASGI) at port 5000.

Note

The prediction API uses app_utils.ModelManager for model lifecycle management and app_utils.NumpyEncoder for JSON serialization of NumPy arrays in detection results.

Model Export

# Export SavedModel after training
python official/vision/train.py \
  --experiment=resnet_imagenet \
  --mode=eval \
  --model_dir=/tmp/trained_model \
  --export_dir=/tmp/saved_model

For TFLite conversion, use ai-edge-litert:

import ai_edge_litert as tfl

converter = tfl.TFLiteConverter.from_saved_model('/tmp/saved_model')
tflite_model = converter.convert()

5. Operations

Training Lifecycle

1. Config loading — YAML → ExperimentConfig dataclass, with Gin and flag overrides
2. Strategy setup — distribute_utils.get_distribution_strategy() creates the distribution context
3. Task creation — task_factory.get_task() instantiates the task from the experiment registry
4. Training loop — Orbit Controller manages train/eval interleaving, checkpointing, and summaries
5. Checkpointing — tf.train.CheckpointManager with configurable max_to_keep
6. Evaluation — Periodic evaluation at validation_interval steps
7. Export — Optional SavedModel/TFLite export at completion

Monitoring with TensorBoard

tensorboard --logdir=/path/to/model_dir

The Orbit Controller writes training and evaluation summaries at summary_interval steps. Custom summary managers (e.g., ImageScalarSummaryManager in official/vision/utils/summary_manager.py) extend the default behavior.

Checkpoint Management

checkpoint_manager = tf.train.CheckpointManager(
    checkpoint,
    directory=model_dir,
    max_to_keep=5,
    checkpoint_name='ckpt')

Orbit's Controller supports:

• Synchronous checkpointing (default)
• Async checkpointing (enable_async_checkpointing=True in vision train.py)
• EMA checkpointing via EMACheckpointing action (saves exponential moving average weights under ema_checkpoints/)

Preemption Recovery

The vision training driver (official/vision/train.py) includes a preemption recovery loop that automatically reconnects when Cloud TPU preemption occurs:

def _run_experiment_with_preemption_recovery(params, model_dir):
    keep_training = True
    while keep_training:
        try:
            # Setup strategy and run training
            ...
        except (tf.errors.UnavailableError, ...):
            # Reconnect and resume from last checkpoint
            continue

Tip

For long-running TPU training jobs, set checkpoint_interval to a reasonable frequency (every 1000–5000 steps). The preemption recovery loop restarts from the last saved checkpoint.

Continuous Evaluation

controller = orbit.Controller(
    global_step=global_step,
    evaluator=evaluator,
    checkpoint_manager=checkpoint_manager)

# Monitor a directory for new checkpoints and evaluate each
controller.evaluate_continuously(
    timeout=3600,          # Max wait time for new checkpoints
    timeout_fn=timeout_fn  # Custom timeout logic
)

Note

evaluate_continuously is designed for asynchronous training/evaluation pipelines. Point it at the model_dir of a running training job and it will automatically evaluate each new checkpoint.

6. Performance

Mixed Precision Training

Mixed precision is automatically enabled for TPU (use_bfloat16=True). For GPU:

runtime:
  mixed_precision_dtype: 'float16'

Or in the legacy detection config:

'architecture': {
    'use_bfloat16': True,  # bfloat16 for TPU
}

Tip

Set steps_per_loop=1000 for TPU training. Running steps inside a tf.function dramatically increases TPU utilization compared to per-step Python overhead.

XLA Compilation

python official/legacy/detection/main.py --enable_xla=True ...

XLA (Accelerated Linear Algebra) compiles TF graphs for optimized execution. Benefits vary by model; object detection models with dynamic shapes may see limited gains.

Gradient Optimization

The official/modeling/grad_utils.py module provides:

• _filter_grads — Removes None gradients (for frozen layers) and raises ValueError if no gradients remain
• _filter_and_allreduce_gradients — Performs distributed all-reduce in configurable precision (float32 or float16) with optional packing (bytes_per_pack)

from official.modeling.grad_utils import _filter_and_allreduce_gradients

allreduced_grads, variables = _filter_and_allreduce_gradients(
    grads_and_vars,
    allreduce_precision="float16",  # Reduce communication overhead
    bytes_per_pack=128 * 1024 * 1024  # 128 MB packs
)

Frozen Variable Optimization

For transfer learning, freeze early layers to reduce computation:

# Legacy detection: freeze ResNet conv1 + conv2_x
RESNET_FROZEN_VAR_PREFIX = r'(resnet\d+)\/(conv2d(|_([1-9]|10))|batch_normalization(|_([1-9]|10)))\/'

Set train.frozen_variable_prefix to this regex in the detection config.

7. Security

Input Validation

• Image uploads (waste identification API): Validates image format with PIL, returns HTTP 400 for corrupt images and HTTP 422 for wrong format
• Config validation: ParamsDict enforces type constraints and restrictions; params.validate() checks inter-parameter consistency before params.lock() makes config immutable
• File type validation: pick_dataset_fn() only accepts tfrecord and tfrecord_compressed; raises ValueError for unknown types

Distribution Strategy Safety

• GPU count validation: negative num_gpus raises ValueError
• Strategy/GPU mismatch: off strategy with num_gpus > 0 raises ValueError
• one_device with multiple GPUs raises ValueError
• Multi-worker without task_index raises ValueError

Checkpoint Security

• Checkpoints contain model weights as TensorFlow variables; no executable code in standard tf.train.Checkpoint format
• SavedModel format includes serialized graph — load only from trusted sources

Caution

Pre-trained models and checkpoints should only be loaded from trusted sources (official Google Cloud Storage buckets, TF Hub). A malicious SavedModel can execute arbitrary TensorFlow ops during loading. Verify checkpoint provenance before use in production.

Network Exposure

• The waste identification prediction API (app.py) binds to 0.0.0.0:5000 — place behind a reverse proxy in production
• No authentication or rate limiting built into the prediction API
• No sensitive data stored in model artifacts

8. Developer Internals

Repository Structure

models/
├── tensorflow_models/              # pip package entry: import tensorflow_models as tfm
│   ├── __init__.py                 # Re-exports: nlp, vision, uplift, core, optimization
│   ├── nlp.py                      # tfm.nlp namespace
│   ├── vision.py                   # tfm.vision namespace
│   └── uplift.py                   # tfm.uplift namespace
├── official/                       # Production models (Google-maintained)
│   ├── core/                       # Framework core (27 files, 213K chars)
│   │   ├── config_definitions.py   # ExperimentConfig, RuntimeConfig, TrainerConfig
│   │   ├── train_lib.py            # run_experiment() orchestrator
│   │   ├── train_utils.py          # Config serialization, strategy utilities
│   │   ├── base_task.py            # Abstract base for all tasks
│   │   ├── base_trainer.py         # Base trainer with Orbit integration
│   │   ├── task_factory.py         # get_task() from registry
│   │   ├── exp_factory.py          # Experiment config factory
│   │   ├── registry.py             # Global model/task/config registry
│   │   ├── input_reader.py         # Unified data pipeline
│   │   ├── export_base.py          # SavedModel export utilities
│   │   ├── actions.py              # PruningAction, EMACheckpointing
│   │   ├── file_writers.py         # Output file utilities
│   │   ├── savedmodel_checkpoint_manager.py  # SavedModel-aware checkpoint management
│   │   ├── tf_example_builder.py   # TFExample construction helpers
│   │   └── tf_example_feature_key.py  # Standard feature key constants
│   ├── modeling/                    # Shared modeling utilities (6 files, 28K chars)
│   │   ├── hyperparams/            # ParamsDict, Config, typed parameters
│   │   │   └── params_dict.py      # Hierarchical dict with validation + locking
│   │   ├── optimization/           # Learning rate schedules, EMA, optimizers
│   │   ├── grad_utils.py           # Gradient filtering, all-reduce
│   │   ├── tf_utils.py             # TF utility functions
│   │   ├── performance.py          # Mixed precision, XLA configuration
│   │   └── multitask/              # Multi-task learning utilities
│   ├── common/                      # Shared infrastructure (7 files, 24K chars)
│   │   ├── distribute_utils.py     # Distribution strategy factory
│   │   ├── dataset_fn.py           # Dataset type picker
│   │   └── flags.py                # Common flag registration
│   ├── utils/                       # Utilities
│   │   ├── flags/                   # Flag definitions (_base.py, _perf.py, core.py)
│   │   │   └── _base.py            # define_base() registers standard CLI flags
│   │   ├── testing/                 # Test utilities and mock tasks
│   │   └── misc/                    # Keras utilities
│   ├── vision/                      # Vision models (4 core files + layers/configs/tasks)
│   │   ├── train.py                 # Training driver with preemption recovery
│   │   ├── registry_imports.py      # Registry side-effect imports
│   │   ├── configs/experiments/     # YAML configs by task type
│   │   ├── modeling/                # Backbones, heads, decoders, layers
│   │   ├── tasks/                   # Image classification, detection, segmentation tasks
│   │   ├── dataloaders/             # Data pipeline implementations
│   │   ├── losses/                  # Loss functions
│   │   ├── ops/                     # Vision-specific ops
│   │   ├── evaluation/              # Evaluation metrics
│   │   ├── serving/                 # Serving utilities
│   │   └── utils/                   # Summary manager, object detection utils
│   ├── nlp/                         # NLP models (5 core files + modeling/configs/tasks)
│   │   ├── optimization.py          # WarmUp LR schedule, optimizer factory
│   │   ├── continuous_finetune_lib.py  # Continuous fine-tuning utilities
│   │   ├── modeling/                # BERT, encoders, networks, layers (19 model files)
│   │   ├── configs/                 # NLP experiment configs
│   │   ├── tasks/                   # NLP task implementations
│   │   ├── finetuning/              # GLUE, SQuAD, SuperGLUE fine-tuning
│   │   └── serving/                 # NLP serving utilities
│   ├── recommendation/              # Recommendation models (15 files, 145K chars)
│   │   ├── ranking/                 # Ranking models and data preprocessing
│   │   └── uplift/                  # Uplift modeling (two-tower, treatment effects)
│   ├── legacy/                      # Deprecated models (benchmarking only)
│   │   ├── detection/               # RetinaNet, Mask R-CNN, ShapeMask
│   │   ├── bert/                    # Legacy BERT implementation
│   │   ├── albert/                  # ALBERT (deprecated)
│   │   ├── xlnet/                   # XLNet
│   │   ├── image_classification/    # ResNet, EfficientNet classifier trainer
│   │   └── transformer/             # Legacy transformer
│   ├── projects/                    # 40+ domain-specific projects
│   │   ├── yolo/                    # YOLOv4, YOLOv7, Scaled-YOLO
│   │   ├── centernet/               # CenterNet detection
│   │   ├── detr/                    # DETR detection transformer
│   │   ├── simclr/                  # Self-supervised contrastive learning
│   │   ├── movinet/                 # Mobile video networks
│   │   ├── perceiver/               # Perceiver models
│   │   ├── longformer/              # Longformer attention
│   │   ├── bigbird/                 # BigBird sparse attention
│   │   ├── mae/                     # Masked autoencoders
│   │   ├── maxvit/                  # MaxViT vision transformer
│   │   ├── pointpillars/            # 3D object detection
│   │   ├── panoptic/                # Panoptic segmentation
│   │   ├── maskconver/              # MaskConver
│   │   ├── pix2seq/                 # Pix2Seq
│   │   ├── qat/                     # Quantization-aware training (vision + NLP)
│   │   ├── pruning/                 # Model pruning
│   │   ├── edgetpu/                 # Edge TPU optimized models
│   │   ├── videoglue/               # Video understanding
│   │   ├── waste_identification_ml/ # CircularNet waste detection + prediction API
│   │   └── ...                      # basnet, fffner, unified_detector, etc.
│   └── pip_package/                 # setup.py for PyPI distribution
├── orbit/                           # Training loop library
│   ├── __init__.py                  # Public API: Controller, runners, actions
│   ├── controller.py                # Controller class (train/eval outer loop)
│   ├── runner.py                    # AbstractTrainer, AbstractEvaluator interfaces
│   ├── standard_runner.py           # StandardTrainer, StandardEvaluator
│   ├── actions/                     # Built-in actions (8 files, 38K chars)
│   │   ├── __init__.py              # Export, NewBestMetric, ConditionalAction
│   │   └── ...                      # Action implementations
│   ├── utils/                       # Summary managers, utilities (9 files, 33K chars)
│   └── examples/                    # Example trainer implementations
├── research/                        # Paper implementations (community-maintained)
│   ├── object_detection/            # TF OD API (44 core, 62 utils, 81 models)
│   │   ├── protos/                  # 34 protobuf definitions
│   │   ├── core/                    # Core detection abstractions
│   │   ├── models/                  # SSD, Faster R-CNN, and Keras model implementations
│   │   ├── utils/                   # Visualization, evaluation, ops
│   │   ├── builders/                # Builder pattern for constructing models from configs
│   │   ├── meta_architectures/      # Meta-architecture abstractions
│   │   └── dataset_tools/           # Data conversion utilities
│   ├── deeplab/                     # DeepLab segmentation (12 core files + core/)
│   │   ├── core/                    # Xception, ResNet backbones, dense prediction
│   │   ├── utils/                   # Training, evaluation utilities
│   │   └── datasets/                # Dataset definitions
│   ├── slim/                        # TF-Slim model zoo
│   │   ├── nets/                    # MobileNet, NASNet, network definitions
│   │   └── preprocessing/           # Image preprocessing pipelines
│   ├── delf/                        # DEep Local Features
│   │   └── delf/python/examples/    # CLI tools: extract_boxes, extract_features, match_images
│   ├── audioset/                    # VGGish, YAMNet audio classifiers
│   ├── adversarial_text/            # Adversarial training for text
│   ├── seq_flow_lite/               # On-device sequence models
│   ├── autoaugment/                 # Data augmentation policies
│   ├── deep_speech/                 # Speech recognition
│   ├── cognitive_planning/          # Cognitive planning research
│   ├── lfads/                       # Latent Factor Analysis via Dynamical Systems
│   ├── pcl_rl/                      # Path Consistency Learning (RL)
│   ├── vid2depth/                   # Video-to-depth estimation
│   ├── rebar/                       # REBAR gradient estimator
│   ├── efficient-hrl/               # Efficient hierarchical RL
│   ├── lstm_object_detection/       # LSTM-based video object detection
│   ├── marco/                       # Automated image analysis
│   ├── cvt_text/                    # Cross-view training for text
│   └── nst_blogpost/                # Neural style transfer
├── docs/                            # Public documentation (mirrored to tensorflow.org/tfmodels)
│   ├── orbit/                       # Orbit documentation
│   ├── vision/                      # Vision documentation
│   └── nlp/                         # NLP documentation
└── community/                       # Community contributions

Architecture Deep Dive

Registry and Factory Pattern

The official/core/registry.py module implements a global registry for experiments, tasks, and models. Decorators register implementations:

from official.core import exp_factory

@exp_factory.register_config_factory('resnet_imagenet')
def resnet_imagenet() -> cfg.ExperimentConfig:
    return cfg.ExperimentConfig(
        task=ResNetImageNetConfig(),
        trainer=TrainerConfig(train_steps=90000))

task_factory.get_task(params.task) looks up the registered task class and instantiates it. This pattern allows adding new experiments without modifying framework code.

Orbit Controller Internals

The Controller class (orbit/controller.py) implements four outer loop patterns:

1. train(steps) — Calls trainer.train(steps_per_loop) in a loop, running train actions after each inner loop, saving checkpoints at intervals
2. evaluate(steps) — Runs evaluator.evaluate(steps) once, runs eval actions
3. train_and_evaluate(train_steps, eval_steps, eval_interval) — Interleaves training and evaluation
4. evaluate_continuously(timeout) — Monitors a checkpoint directory, evaluates each new checkpoint

The Action type is Callable[[runner.Output], None] — any callable that accepts training/eval output. Built-in actions include:

• PruningAction — Updates pruning step counters and logs pruning summaries
• EMACheckpointing — Swaps model weights with EMA averages and saves checkpoints

Configuration System Internals

ParamsDict (official/modeling/hyperparams/params_dict.py):

• Hierarchical dictionary with dot-notation access
• Supports override() with strict mode (key must exist) or relaxed mode
• validate() checks parameter restrictions (e.g., mutually exclusive options)
• lock() makes the config immutable after validation
• Used by legacy detection models

ExperimentConfig (official/core/config_definitions.py):

• Dataclass-based config with type annotations
• Includes RuntimeConfig, TrainerConfig, and task-specific sub-configs
• YAML serializable/deserializable
• Used by all modern official/ models

Data Pipeline Architecture

1. official/core/input_reader.py — Unified input reader supporting TFRecord and TFDS
2. official/common/dataset_fn.py — pick_dataset_fn() returns the appropriate tf.data.Dataset constructor based on file type
3. Task-specific dataloaders in official/vision/dataloaders/ and official/nlp/data/
4. Research models use their own data loading (e.g., research/object_detection/inputs/)

NLP Optimization

official/nlp/optimization.py provides:

• WarmUp — Linear warmup learning rate schedule wrapping any base schedule
• Optimizer factory supporting Adam, AdamW, LAMB, and SGD with configurable weight decay, warmup, and polynomial/cosine decay

class WarmUp(tf_keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, initial_learning_rate, decay_schedule_fn,
                 warmup_steps, power=1.0):
        # Linear warmup from 0 to initial_learning_rate over warmup_steps
        # Then follows decay_schedule_fn

Module Reference

official/core (27 files, 213K chars)

orbit/ (6 core files + 8 actions + 9 utils, 155K chars total)

official/modeling (6 files, 28K chars)

official/common (7 files, 24K chars)

research/object_detection (44 core + 62 utils + 81 models + 34 protos)

research/deeplab (12 core + 18 core files, 402K chars)

Code Patterns and Conventions

Import Pattern

import tensorflow as tf, tf_keras  # Dual import for TF 2.16+ compatibility

All files use Apache License 2.0 with Google copyright headers.

Registry Pattern

# Registration
@exp_factory.register_config_factory('my_experiment')
def my_experiment():
    return ExperimentConfig(...)

# Usage
config = exp_factory.get_exp_config('my_experiment')

Config Override Chain

Default Config → YAML File → Gin Bindings → CLI --params_override → params.lock()

Error Handling Pattern

The codebase uses explicit ValueError for configuration errors, NotImplementedError for abstract interfaces, and OverflowError for numerical instability (NaN loss). Of the 694 error messages extracted, the majority (>500) are ValueError for invalid parameter combinations.

Testing Pattern

Tests use tf.test.TestCase and parameterized.TestCase (from absl-py). The test files (*_test.py) are excluded from the pip package but present in the source tree.

Testing Strategy

# Run all tests for a module
python -m pytest official/core/ -v

# Run a specific test
python -m pytest official/common/distribute_utils_test.py -v

Test utilities in official/utils/testing/ include mock_task.py for creating lightweight task mocks during testing.

Potential Improvements

Based on source analysis:

1. Legacy code removal — The official/legacy/ directory contains deprecated implementations (BERT, ALBERT, image classification, detection) that duplicate functionality in official/nlp/ and official/vision/. Consider archiving or removing once all benchmarks migrate.

2. Research model maintenance — Several research models depend on outdated TF APIs and may not work with the latest TensorFlow (e.g., seq_flow_lite build issues, GitHub issue #10737).

3. Prediction API hardening — The waste identification FastAPI server lacks authentication, rate limiting, input size validation, and health check endpoints. Production deployment needs a reverse proxy and request throttling.

4. Config system unification — Two config systems coexist: ParamsDict (legacy detection) and ExperimentConfig (modern). Migrating legacy detection to ExperimentConfig would reduce maintenance burden.

5. Protobuf compilation step — The object detection API requires manual protobuf compilation (protoc object_detection/protos/*.proto --python_out=.). Pre-compiled Python files or build automation would improve onboarding.

6. Checkpoint URL freshness — Multiple GitHub issues report 403/AccessDenied errors for hosted checkpoints. A CI job validating checkpoint URL accessibility would catch these proactively.

7. Python 3.13 compatibility — GitHub issue #13587 reports installation issues on Python 3.13. Dependency resolution and wheel availability need tracking.

Contributing

1. Fork the repository on GitHub
2. Set up development environment:

   git clone https://github.com/YOUR_USER/models.git
   cd models
   pip install -r official/requirements.txt
   export PYTHONPATH="$PYTHONPATH:$(pwd)"
   

3. Create a feature branch
4. Follow the existing code style: Apache 2.0 headers, tf_keras imports, tf.test.TestCase for tests
5. Run relevant tests:

   python -m pytest official/core/ -v
   

6. Submit a pull request against the master branch

Important

All code contributions must include the Apache License 2.0 header. Test files should follow the naming convention *_test.py and use tf.test.TestCase as the base class.

9. Integration Points

TensorFlow Hub

Models can be loaded from TF Hub instead of raw checkpoints:

import tensorflow_hub as hub

# BERT models
bert = hub.load("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/4")

# Use with training scripts
python run_classifier.py \
  --hub_module_url=https://tfhub.dev/tensorflow/bert_en_uncased_L-24_H-1024_A-16/1

Available BERT hub modules:

• bert_en_uncased_L-12_H-768_A-12 — BERT-Base Uncased (110M params)
• bert_en_uncased_L-24_H-1024_A-16 — BERT-Large Uncased (340M params)
• bert_en_cased_L-12_H-768_A-12 — BERT-Base Cased (110M params)
• bert_en_cased_L-24_H-1024_A-16 — BERT-Large Cased (340M params)
• bert_en_wwm_uncased_L-24_H-1024_A-16 — BERT-Large WWM Uncased (340M params)
• bert_en_wwm_cased_L-24_H-1024_A-16 — BERT-Large WWM Cased (340M params)
• bert_multi_cased_L-12_H-768_A-12 — Multilingual BERT (104 languages, 110M params)
• bert_zh_L-12_H-768_A-12 — Chinese BERT (110M params)

TensorFlow Datasets (TFDS)

The classifier_trainer.py supports TFDS for ImageNet and other datasets:

task:
  train_data:
    tfds_name: 'imagenet2012'
    tfds_split: 'train'

TensorFlow Model Optimization Toolkit

Pruning and quantization-aware training (QAT) are integrated via official/projects/qat/ and official/projects/pruning/:

import tensorflow_model_optimization as tfmot

# PruningAction updates pruning step after each training loop
pruning_action = PruningAction(
    export_dir='/tmp/pruning_summaries',
    model=model,
    optimizer=optimizer)

TensorFlow Lite

Models can be converted to TFLite for on-device inference:

converter = tf.lite.TFLiteConverter.from_saved_model('/tmp/saved_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

The ai-edge-litert package (required dependency) provides the runtime. Research projects like seq_flow_lite include custom TFLite ops (research/seq_flow_lite/tflite_ops/, research/seq_flow_lite/tf_ops/).

Gin-Config

Many experiments support Gin for fine-grained parameterization:

python official/vision/train.py \
  --experiment=resnet_imagenet \
  --gin_file=custom_overrides.gin \
  --gin_params="TrainerConfig.train_steps=50000"

Google Cloud Integration

• Cloud TPU: Native support via tf.distribute.TPUStrategy
• Cloud Storage: All data/checkpoint paths support gs:// URIs
• AI Platform: Research object detection includes Docker configs for AI Platform training (research/object_detection/dockerfiles/tf2_ai_platform/)