TPU-Model-Training

Model files, training and evaluation scripting, and validation sets in support of open-source embedded TPUs for picosatellites.

This repository characterizes how MobileNetV2-style convolutional networks trade off accuracy, size, and latency when quantized and compiled for Google's Edge TPU (Coral), with an eye toward deployment on resource-constrained, radiation-tolerant flight hardware. It contains:

• A configurable architecture-search pipeline that trains a grid of MobileNetV2 variants (width multiplier × body depth) on ImageNet.
• Post-training INT8 quantization and Edge TPU compilation for every trained model.
• CPU-based evaluation harnesses for both image classification (ImageNet) and object detection (COCO), used to benchmark both the custom sweep and a set of pre-made reference models (stock MobileNet/EfficientNet-EdgeTPU/Inception/SSD variants).
• A standalone transfer-learning example (MobileNetV2 → CIFAR-100) showing the same train → quantize → export flow on a smaller dataset.

Repository layout

Path	Contents
src/model_sweeps/	Core pipeline: config generation, training, TFLite/Edge TPU export, dataset utilities
src/detection_eval/	CPU evaluation of TFLite models on ImageNet (classification) and COCO (detection)
src/retraining/	Standalone MobileNetV2 → CIFAR-100 transfer-learning example
src/scripts/	Ad hoc analysis scripts
data/	Datasets (ImageNet, COCO) and reference/baseline model files
results/	Training logs, exported models, evaluation JSON, and generated plots
mobilenetv2_cifar100_savedmodel/	Example SavedModel export produced by src/retraining/train_image_class

See the README in each subdirectory for details specific to that part of the pipeline.

Pipeline overview

model_config.py            model_training.py             export_TPU.py
generates a grid of   -->  trains each config on    -->  quantizes to INT8 TFLite
(alpha, depth) configs     ImageNet (60/20/20 split)      and compiles for Edge TPU
        |                          |                              |
        v                          v                              v
sweep_configs_*.json     results/*/models/*.keras     results/*/tflite_models/*.tflite
                          results/*/training_logs/*.csv  results/*/edgeTPU/*_edgetpu.tflite

Evaluation is a separate pass over the exported .tflite files:

src/detection_eval/eval_imagenet_cpu.py  -->  results/**/*_eval.json  (top-1/5, precision, recall, latency)
src/detection_eval/eval_coco_cpu.py      -->  data/models/**/*_coco_eval.json  (COCO mAP)
src/detection_eval/eval_imagenet_224.py  -->  results/plots/*.png  (parses the eval JSON/logs above into sweep plots)

alpha is the MobileNetV2 width multiplier (0.25–1.5); depth/depth_repeats is the number of repeated stride-1 blocks in the network body (2–12, standard MobileNetV2 uses 5). Model IDs follow the pattern Grid_A<alpha>_D<depth>, e.g. Grid_A1.0_D08.

Setup

The scripts target a Linux workstation with a ROCm or CUDA GPU (see the Docker invocation documented at the top of src/retraining/train_image_class for a known-good ROCm container). There is currently no pinned requirements.txt; the pipeline depends on:

tensorflow (with GPU support), pycoral, pycocotools, opencv-python, Pillow,
numpy, pandas, matplotlib, scikit-learn, tqdm

Quantized models are compiled for Edge TPU with the edgetpu_compiler command-line tool (see the Coral documentation), which must be installed separately and available on PATH.

Every entry point resolves the repo root via utils.get_repo_root() (src/model_sweeps/utils.py), which checks the CORAL_TRAIN_REPO environment variable first, then a short list of common clone locations (~/TPU-Model-Training, ~/Dev/repos/TPU-Model-Training, /app, ~/Documents/TPU-Model-Training). Set CORAL_TRAIN_REPO if you clone elsewhere.

Datasets (ImageNet, COCO, CIFAR-100) are not checked into the repository (see .gitignore) and must be downloaded separately — see data/README.md.

Quick start

# 1. Generate the sweep grid (or use the committed sweep_configs_example.json)
python src/model_sweeps/model_config.py

# 2. Train every config in the grid
python src/model_sweeps/model_training.py

# 3. Quantize (INT8) and compile every trained model for Edge TPU
python src/model_sweeps/export_TPU.py

# 4. Evaluate the exported TFLite models on a held-out ImageNet split
python src/detection_eval/eval_imagenet_cpu.py --models results/model_sweeps_new/tflite_models --dataset data/imagenet/test_224_20

# 5. Turn the eval JSON into comparison plots
python src/detection_eval/eval_imagenet_224.py