# 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

```text
┌─────────────────────────────────────────────────────────────────────┐
│                     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

| Component | Path | Role |
|---|---|---|
| `official/core` | `official/core/` | Registry, experiment factory, task base classes, training orchestration |
| `official/modeling` | `official/modeling/` | Shared modeling utilities: hyperparameters, optimization, gradient tools |
| `official/common` | `official/common/` | Distribution strategy helpers, dataset utilities, common flags |
| `official/vision` | `official/vision/` | Vision models: classification, detection, segmentation, video |
| `official/nlp` | `official/nlp/` | NLP models: BERT, encoders, fine-tuning pipelines |
| `official/recommendation` | `official/recommendation/` | Recommendation models, uplift modeling |
| `official/legacy` | `official/legacy/` | Deprecated models kept for benchmarking (BERT, detection, image classification) |
| `official/projects` | `official/projects/` | Domain-specific projects (40+ sub-projects) |
| `orbit` | `orbit/` | Training loop library: Controller, runners, actions, utilities |
| `tensorflow_models` | `tensorflow_models/` | Public pip package entry point (`import tensorflow_models as tfm`) |
| `research/object_detection` | `research/object_detection/` | TF Object Detection API (SSD, Faster R-CNN, etc.) |
| `research/deeplab` | `research/deeplab/` | DeepLab semantic segmentation |
| `research/slim` | `research/slim/` | TF-Slim model library |
| `research/delf` | `research/delf/` | DEep Local Features for image retrieval |
| `research/audioset` | `research/audioset/` | VGGish and YAMNet audio classifiers |

### Technology Stack

| Layer | Technology |
|---|---|
| Runtime | Python 3.7+, TensorFlow 2.20.0 |
| ML Framework | `tf_keras` 2.16+ (standalone Keras) |
| Distribution | `tf.distribute` (Mirrored, TPU, MultiWorker, ParameterServer) |
| Configuration | YAML (`ExperimentConfig` dataclasses), Gin-config, `absl-py` flags |
| Data Pipeline | `tf.data`, TFRecord, TFDS |
| Serialization | Protocol Buffers (object detection protos), SavedModel |
| Package | `tf-models-official` on PyPI, `setuptools` |
| CV Dependencies | `opencv-python-headless`, `Pillow`, `pycocotools`, `Cython` |
| NLP Dependencies | `sentencepiece`, `seqeval`, `sacrebleu`, `tensorflow-text` |
| Optimization | `tensorflow-model-optimization` (pruning, quantization) |
| Serving | TFLite (`ai-edge-litert`), SavedModel, FastAPI (waste_identification_ml) |

> [!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

| Requirement | Development | Production (GPU) | Production (TPU) |
|---|---|---|---|
| Python | 3.7+ | 3.10+ | 3.10+ |
| TensorFlow | 2.20.0 | 2.20.0 | tf-nightly or 2.20.0 |
| CPU | Any x86_64 | 8+ cores | N/A (Cloud VM) |
| RAM | 8 GB | 32 GB+ | 16 GB+ (VM) |
| GPU | Optional | NVIDIA A100/V100/T4 with CUDA 12.x | N/A |
| Storage | 2 GB (code + deps) | 100 GB+ (datasets) | GCS bucket |
| Network | Internet for pip | High bandwidth for data loading | GCP internal |

### Software Dependencies (from `official/requirements.txt`)

| Package | Required Version | Purpose |
|---|---|---|
| `tensorflow` | ~=2.20.0 | Core ML framework |
| `tensorflow-text` | ~=2.20.0 | NLP text processing |
| `tf-keras` | >=2.16.0 | Keras API |
| `tensorflow-hub` | >=0.6.0 | Pre-trained model loading |
| `tensorflow-model-optimization` | >=0.4.1 | Pruning, QAT |
| `tensorflow-datasets` | any | Dataset loading |
| `gin-config` | any | Experiment configuration |
| `tf_slim` | >=1.1.0 | Slim model definitions |
| `numpy` | >=1.20 | Numerical computing |
| `scipy` | >=0.19.1 | Scientific computing |
| `pandas` | >=0.22.0 | Data processing |
| `Pillow` | any | Image processing |
| `opencv-python-headless` | any | Computer vision operations |
| `pycocotools` | any | COCO evaluation metrics |
| `pyyaml` | >=6.0.0 | YAML config parsing |
| `kaggle` | >=1.3.9 | Dataset downloads |
| `ai-edge-litert` | >=1.0.1 | TFLite runtime |
| `sentencepiece` | any | NLP tokenization |
| `seqeval` | any | Sequence labeling metrics |
| `sacrebleu` | any | MT evaluation metrics |
| `immutabledict` | any | Immutable dictionaries |

---

## 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

| File/Directory | Format | Purpose |
|---|---|---|
| `official/vision/configs/experiments/*.yaml` | YAML | Vision task experiment configs |
| `official/nlp/configs/experiments/*.yaml` | YAML | NLP task experiment configs |
| `official/projects/*/configs/experiments/*.yaml` | YAML | Per-project experiment configs |
| `official/legacy/detection/configs/base_config.py` | Python dict | Legacy detection base configuration |
| `official/legacy/detection/configs/*_config.py` | Python dict | Per-model detection configs |

### Environment Variables

| Variable | Type | Default | Description | Required? |
|---|---|---|---|---|
| `TF_CONFIG` | JSON string | `{}` | Multi-worker cluster specification. Read by `official/common/distribute_utils.py:217` to configure multi-worker training. Must include `cluster` and `task` keys. | Only for multi-worker |
| `PYTHONPATH` | Path list | — | Must include repository root for source installs | Source installs only |
| `CUDA_VISIBLE_DEVICES` | Comma-separated ints | all GPUs | Restrict visible GPUs | No |

### CLI Flags — Common Training Flags

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

| Flag | Short | Type | Default | Description |
|---|---|---|---|---|
| `--data_dir` | `-dd` | `str` | `/tmp` | Location of input data |
| `--model_dir` | `-md` | `str` | `/tmp` | Location for model checkpoints |
| `--clean` | — | `bool` | `False` | Remove `model_dir` before training |
| `--train_epochs` | `-te` | `int` | `1` | Number of training epochs |
| `--epochs_between_evals` | `-ebe` | `int` | `1` | Epochs between evaluations |
| `--stop_threshold` | `-st` | `float` | `None` | Stop training when metric exceeds this |
| `--batch_size` | `-bs` | `int` | `32` | Global batch size (split across devices) |
| `--num_gpus` | `-ng` | `int` | `1` | GPUs per worker |
| `--distribution_strategy` | `-ds` | `str` | `mirrored` | Distribution strategy (`off`, `one_device`, `mirrored`, `parameter_server`, `multi_worker_mirrored`, `tpu`) |
| `--run_eagerly` | — | `bool` | `False` | Run ops eagerly (for debugging) |
| `--hooks` | `-hk` | `list` | `LoggingTensorHook` | Training hooks |
| `--export_dir` | `-ed` | `str` | `None` | Export SavedModel after training |
| `--log_steps` | — | `int` | — | Log metrics every N steps |

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

| Flag | Type | Default | Description |
|---|---|---|---|
| `--experiment` | `str` | (required) | Registered experiment name |
| `--config_file` | `str` | — | YAML config path |
| `--params_override` | `str` | — | Comma-separated key=value overrides |
| `--mode` | `str` | `train` | `train`, `eval`, `train_and_eval`, `continuous_eval` |
| `--enable_async_checkpointing` | `bool` | `True` | Enable async checkpoint saving |

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

| Flag | Type | Default | Description |
|---|---|---|---|
| `--mode` | `str` | `train` | `train`, `eval`, or `eval_once` |
| `--model` | `str` | `retinanet` | `retinanet`, `mask_rcnn`, `olnmask`, `shapemask` |
| `--enable_xla` | `bool` | `False` | Enable XLA compilation for GPU |
| `--training_file_pattern` | `str` | `None` | Training data glob pattern |
| `--eval_file_pattern` | `str` | `None` | Evaluation data glob pattern |
| `--checkpoint_path` | `str` | `None` | Checkpoint for `eval_once` |

### Legacy Detection Base Configuration

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

| Section | Key | Default | Description |
|---|---|---|---|
| `train` | `batch_size` | `64` | Training batch size |
| `train` | `total_steps` | `22500` | Total training steps |
| `train` | `iterations_per_loop` | `100` | Steps per training loop iteration |
| `train.optimizer` | `type` | `momentum` | Optimizer type |
| `train.optimizer` | `momentum` | `0.9` | Momentum value |
| `train.optimizer` | `nesterov` | `True` | Use Nesterov momentum |
| `train.learning_rate` | `type` | `step` | LR schedule type |
| `train.learning_rate` | `warmup_learning_rate` | `0.0067` | Warmup LR |
| `train.learning_rate` | `warmup_steps` | `500` | Warmup duration |
| `train.learning_rate` | `init_learning_rate` | `0.08` | Initial learning rate |
| `train` | `l2_weight_decay` | `0.0001` | L2 regularization |
| `train` | `gradient_clip_norm` | `0.0` | Gradient clipping (0 = disabled) |
| `eval` | `batch_size` | `8` | Eval batch size |
| `eval` | `eval_samples` | `5000` | Number of eval samples |
| `eval` | `min_eval_interval` | `180` | Minimum seconds between evals |
| `architecture` | `backbone` | `resnet` | Backbone network (`resnet`, `spinenet`) |
| `architecture` | `min_level` | `3` | Minimum FPN level |
| `architecture` | `max_level` | `7` | Maximum FPN level |
| `architecture` | `multilevel_features` | `fpn` | Feature extraction (`fpn`, `identity`) |
| `architecture` | `use_bfloat16` | `True` | Use bfloat16 mixed precision |
| `architecture` | `num_classes` | `91` | Number of classes (including background) |
| `resnet` | `resnet_depth` | `50` | ResNet depth |
| `fpn` | `fpn_feat_dims` | `256` | FPN feature dimensions |
| `postprocess` | `nms_iou_threshold` | `0.5` | NMS IoU threshold |
| `postprocess` | `score_threshold` | `0.05` | Detection score threshold |
| `postprocess` | `max_total_size` | `100` | Max detections per image |

> [!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

```python
# 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)
)
```

| Strategy | `all_reduce_alg` options | Use case |
|---|---|---|
| `mirrored` | `nccl`, `hierarchical_copy` | Single host, multiple GPUs |
| `multi_worker_mirrored` | `ring`, `nccl` | Multiple hosts with GPUs |
| `tpu` | N/A | Google Cloud TPU |
| `one_device` | N/A | Single GPU/CPU (debugging) |
| `parameter_server` | N/A | Async training with parameter servers |

> [!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:

| Package | Version Pattern | TF Dependency |
|---|---|---|
| `tf-models-official` | `2.20.0` | `tensorflow~=2.20.0` |
| `tf-models-nightly` | `2.21.0.dev{YYYYMMDD}` | `tf-nightly` |

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

### Building from Source

```bash
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

```bash
# 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

```bash
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:

```dockerfile
# 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

```python
# 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`:

```python
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

```bash
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

```python
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:

```python
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

```python
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:

```yaml
runtime:
  mixed_precision_dtype: 'float16'
```

Or in the legacy detection config:

```python
'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

```bash
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`)

```python
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:

```python
# 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

```text
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:

```python
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

```python
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)

| Module | Purpose |
|---|---|
| `config_definitions.py` | `ExperimentConfig`, `RuntimeConfig`, `TrainerConfig` dataclasses |
| `train_lib.py` | `run_experiment()` — main orchestrator for training with Orbit |
| `train_utils.py` | Config serialization, strategy utilities, write_json_summary |
| `base_task.py` | Abstract `Task` base class defining build_model, build_losses, train_step, validation_step |
| `base_trainer.py` | `BaseTrainer` integrating Task with Orbit's AbstractTrainer |
| `task_factory.py` | `get_task()` — instantiate task from config using registry |
| `exp_factory.py` | `get_exp_config()` — instantiate experiment config from registry |
| `registry.py` | Decorator-based global registry for experiments, tasks, models |
| `input_reader.py` | Unified `InputReader` for TFRecord/TFDS data pipelines |
| `export_base.py` | SavedModel export base classes and utilities |
| `actions.py` | `PruningAction`, `EMACheckpointing` — Orbit train/eval actions |
| `savedmodel_checkpoint_manager.py` | CheckpointManager that also exports SavedModel at each save |
| `tf_example_builder.py` | Helpers for constructing `tf.train.Example` protos |
| `tf_example_feature_key.py` | Standard feature key constants for TFExample |
| `file_writers.py` | File output utilities |

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

| Module | Purpose |
|---|---|
| `controller.py` | `Controller` class — outer training/eval loop management |
| `runner.py` | `AbstractTrainer`, `AbstractEvaluator` — interfaces for inner loops |
| `standard_runner.py` | `StandardTrainer`, `StandardEvaluator` — default implementations with `tf.function` tracing |
| `actions/` | Built-in actions: export, metric tracking, conditional execution |
| `utils/` | `SummaryManager`, `SummaryManagerInterface`, epoch helpers |

#### `official/modeling` (6 files, 28K chars)

| Module | Purpose |
|---|---|
| `grad_utils.py` | Gradient filtering, distributed all-reduce with precision control |
| `hyperparams/params_dict.py` | Hierarchical config with validation, overrides, and locking |
| `optimization/` | Learning rate schedules (warmup, cosine, step, polynomial), EMA optimizer |
| `tf_utils.py` | TensorFlow utility functions |
| `performance.py` | Mixed precision and XLA configuration helpers |

#### `official/common` (7 files, 24K chars)

| Module | Purpose |
|---|---|
| `distribute_utils.py` | `get_distribution_strategy()` factory, `configure_cluster()`, strategy helpers |
| `dataset_fn.py` | `pick_dataset_fn()` — returns TFRecordDataset or compressed variant |
| `flags.py` | Common flag registration and TFM-specific flag definitions |

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

| Component | Files | Purpose |
|---|---|---|
| `core/` | 44 files, 1.2M chars | Core abstractions: box lists, target assigners, preprocessor, post-processing, losses |
| `utils/` | 62 files, 1M chars | Visualization, evaluation, ops, config utils |
| `models/` | 81 files, 736K chars | SSD, Faster R-CNN, CenterNet, and Keras model implementations |
| `protos/` | 34 files, 150K chars | Protocol Buffer definitions for detection pipeline configs |
| `meta_architectures/` | — | SSD, Faster R-CNN, CenterNet meta-architectures |
| `builders/` | — | Builder pattern for constructing models from protobuf configs |
| `dataset_tools/` | — | Data conversion (COCO, PASCAL VOC, DensePose, Context R-CNN) |

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

| Component | Purpose |
|---|---|
| `core/` | Xception backbone (`xception.py`), ResNet-v1-beta, dense prediction cell, preprocessing |
| `core/conv2d_ws.py` | Weight-standardized Conv2D layer |
| `core/dense_prediction_cell.py` | Neural architecture search cell for dense prediction |
| `utils/` | Training utilities, evaluation helpers |
| `datasets/` | Dataset registration (Cityscapes, PASCAL VOC, ADE20K) |

### Code Patterns and Conventions

#### Import Pattern

```python
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

```python
# 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

```text
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

| Test Type | Location | Framework |
|---|---|---|
| Unit tests | Co-located with source (`*_test.py`) | `tf.test.TestCase`, `absl.testing` |
| Integration tests | Per-project test directories | `tf.test.TestCase` |
| Config tests | `configs/*_test.py` | Validates config roundtrip |

```bash
# 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:
   ```bash
   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:
   ```bash
   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:

```python
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:

```yaml
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/`:

```python
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:

```python
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:

```bash
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/`)

---

_Generated May 2026 from tensorflow/models at 971ded9. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._