This notebook presents an advanced TPU-based classification pipeline for the Petals to the Metal competition. The objective is to maximize macro F1 score across 104 imbalanced flower classes using a distributed training setup and a dual-architecture ensemble.
The solution integrates:
- TPU-accelerated training
- TFRecord-based data pipelines
- stochastic regularization
- cyclical learning rate scheduling
- dual-stream architecture (EfficientNet + DenseNet)
- test-time augmentation (TTA)
Each component is designed to address a specific limitation in fine-grained image classification.
The dataset contains 104 botanical classes with:
- strong class imbalance
- subtle inter-class visual differences
- limited samples for certain categories
These properties make the problem sensitive to:
- overfitting
- feature under-representation
- unstable convergence
The pipeline focuses on improving both representation capacity and generalization stability.
The environment is explicitly configured to ensure TPU compatibility:
- disabling conflicting acceleration backends (
JAX_PLATFORMS = cpu) - enabling legacy Keras execution
- installing
tensorflow-tpuandlibtpu
TPU kernels require strict runtime alignment. Misconfiguration prevents proper hardware utilization and can halt execution.
A TPUStrategy is initialized to distribute computation across 8 TPU replicas.
Each replica computes gradients independently, followed by synchronous aggregation:
bfloat16is used for computation
- improves throughput
- reduces memory usage
- enables larger batch sizes without instability
Key parameters:
- image size:
512 × 512 - epochs:
16 - batch size scaled as:
Scaling batch size with TPU replicas ensures efficient utilization while maintaining stable gradient updates.
The dataset is dynamically resolved across multiple Kaggle input paths and GCS fallback.
Sample counts:
- training: 12753
- validation: 3712
- test: 7382
- ensures correct dataset mounting across TPU environments
- defines training loop bounds and evaluation frequency
- TFRecord parsing for efficient I/O
- parallel reads with controlled thread usage
- prefetching using
AUTOTUNE
- JPEG decoding
- normalization:
- random flips (horizontal and vertical)
- brightness perturbation
- contrast and saturation scaling
- validation dataset caching removed to prevent RAM exhaustion
TFRecord pipelines are essential for TPU throughput. Augmentation introduces variability to counter overfitting.
A micro-batch is visualized after augmentation.
- verifies preprocessing correctness
- ensures label alignment
- detects numerical anomalies introduced during augmentation
The learning rate schedule consists of:
- linear warmup phase
- optional sustain phase
- exponential decay
- stabilizes early optimization
- prevents divergence in large-batch training
- enables fine convergence during later epochs
The model combines:
- EfficientNet for spatial efficiency
- DenseNet for dense feature propagation
- EfficientNet captures global structure efficiently
- DenseNet preserves feature reuse and gradient flow
This creates complementary feature representations.
Fine-grained classification benefits from diverse feature extraction pathways.
Training is executed under TPU strategy with:
- distributed batches
- scheduled learning rate
- augmented inputs
The objective is to improve macro-level classification performance across imbalanced classes.
Balanced performance across all classes is critical in this competition setting.
Multiple augmented versions of each test image are generated.
Predictions are aggregated:
TTA reduces prediction variance and improves robustness by averaging across multiple transformed inputs.
The pipeline integrates:
- TPU-based distributed training
- TFRecord-optimized data ingestion
- stochastic data augmentation
- cyclical learning rate scheduling
- dual-stream architecture (EfficientNet + DenseNet)
- test-time augmentation for inference
Each component addresses a specific challenge in fine-grained, imbalanced image classification.
This notebook demonstrates a structured TPU pipeline combining efficient data loading, optimized training dynamics, and complementary model architectures.
Further improvements can focus on:
- class-aware loss functions for imbalance handling
- attention mechanisms for fine-grained localization
- adaptive augmentation strategies
- weighted ensembling of model outputs
These directions can further improve classification performance and robustness.