Framework Design¶

Autoware-ML is built on a modular architecture that separates concerns into key components including configuration, data handling, model definition, training, and deployment, among others. This design makes it easy to add new models and datasets while reusing common infrastructure.

Architecture Overview¶

flowchart TB
    subgraph Optimization [Hyperparameter Tuning]
        Optuna[Optuna]
    end

    subgraph Configuration [Configuration Layer]
        YAML[YAML Configs]
        Optuna --> Hydra[Hydra]
        YAML --> Hydra
    end

    subgraph TrainingPipeline [Training Pipeline]
        RecordTable[Record Table]
        RecordTable --> LightningDataModule[Lightning Data Module]
        LightningDataModule --> Transforms[Transforms]
        Transforms --> Collation[Collation]
        Collation --> BatchTransfer[Batch Transfer]
        BatchTransfer --> Preprocessing[Model Preprocessing]
        Preprocessing --> ForwardPass[Forward Pass]
        ForwardPass --> LossComputation[Loss Computation]
        LossComputation --> BackwardPass[Backward Pass]
    end

    subgraph ModelLayer [Model Definition]
        LightningModule[Lightning Module]
        LightningModule --> Blocks["Blocks"]
        LightningModule --> Optimizers["Optimizers"]
        LightningModule --> Schedulers["Schedulers"]
    end

    subgraph TrainingLoop [Training Orchestration]
        Trainer[Lightning Trainer]
        Trainer --> CustomCallbacks[Custom Callbacks]
        Trainer --> MLflow[MLflow Logger]
        Trainer --> Checkpoints[Checkpoints]
    end

    subgraph Deployment [Deployment Pipeline]
        ModelWeights[Model Weights]
        ModelWeights --> ONNXExport[ONNX Export]
        ONNXExport --> TensorRTEngine[TensorRT Engine]
    end

    Hydra --> LightningDataModule
    Hydra --> LightningModule
    Hydra --> Trainer
    Hydra --> ModelWeights

    style RecordTable fill:#bbdefb,opacity:0.2,stroke:#1976d2
    style LightningDataModule fill:#bbdefb,opacity:0.2,stroke:#1976d2
    style Transforms fill:#bbdefb,opacity:0.2,stroke:#1976d2
    style Collation fill:#bbdefb,opacity:0.2,stroke:#1976d2
    style ModelWeights fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style ONNXExport fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style TensorRTEngine fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Blocks fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style BatchTransfer fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Preprocessing fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style ForwardPass fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style LossComputation fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style BackwardPass fill:#a5d6a7,opacity:0.2,stroke:#05bc23

Legend: CPU operations | GPU operations

Core Components¶

Configuration (Hydra)¶

Everything in Autoware-ML is configured through YAML files processed by Hydra. This enables:

• Hierarchical configs - Inherit from base configs, override specific values
• Runtime overrides - Change any parameter from the command line
• Automatic instantiation - _target_ keys specify Python classes to instantiate via hydra.utils.instantiate()

See Configuration Guide for full details on Hydra syntax.

Data Module¶

The data pipeline reads the record table of a database. A database turns its raw annotations into one Parquet table with a row per lidar frame, see Database Design.

DataModule (extending LightningDataModule) builds the datasets of every split from that table:

• prepare_data() generates the record table of every database the splits read, once
• setup() splits each table by the scenario lists of its database and builds one dataset per dataset source of a split
• every split gets its own dataloader settings (batch size, workers, shuffling, pin_memory)

class DataModule(L.LightningDataModule):
    def __init__(
        self,
        splitter: SplitterInterface,
        sources: Mapping[str, DatasetSource],
        train_dataloader: DataLoaderConfig | None,
        validation_dataloader: DataLoaderConfig | None,
        test_dataloader: DataLoaderConfig | None,
        predict_dataloader: DataLoaderConfig | None,
        train_frame_sampling: FrameSamplingConfig | None,
    ) -> None:
        ...

A DatasetSource names one database, the dataset of every split it serves, and how many times its frames appear in one training epoch (repeat). A split is served by the sources declaring a dataset for it, concatenated in declaration order, so several sources mix their corpora. A dataset reads the database of its source and declares whether the boxes (det3d_supervised) and semantic masks (seg3d_supervised) of that corpus supervise the run. A dataset with a task turned off still serves that task, with an empty box set or ignored point labels, so every sample of the split collates with the others. Every source reads its labels through the taxonomies of the run.

The datasets are instantiated by Hydra without records. In setup() the datamodule calls bind(records) on the dataset of every source, which returns a copy serving the records of that database in the split, and concatenates the bound copies of a split in declaration order.

Temporal models that carry a memory across batches train on a StreamingDataModule, a DataModule that also takes the experiment seed. Every dataloader draws from a GroupStreamingSampler, which joins the scenes, the training split reshuffling them every epoch from seed, and cuts their frames, each scene in time order, into one lane per batch position of every rank. The lanes differ by at most one frame, and a shorter lane ends with a padding slot that repeats its last frame, so every rank runs the same number of batches and every split serves every frame exactly once as a real sample. The batch flags the padding slots in ModelGTBatch.padding_mask: a model leaves them out of its loss, and the logged batch size (num_real_samples()) and the detection metrics leave them out; the detection eval output still lists them, flagged, so every GPU gathers as many frames. A prediction writer drops them by the mask it gets with the batch. Anything that would reorder the frames is refused: frame sampling, a shuffling dataloader and trainer.overfit_batches (shorten a run with limit_train_batches / limit_val_batches instead). A model config switches to it by composing /datamodule/streaming@datamodule: default_streaming_datamodule after its database datamodule.

Dataset¶

BaseDataset turns one record into a ModelGTSample and runs the transforms on it:

class BaseDataset(Dataset):
    def __getitem__(self, index: int) -> ModelGTSample:
        model_gt_sample = self.get_data_sample(index)
        return self.apply_transforms(model_gt_sample)

    @abstractmethod
    def get_data_sample(self, index: int) -> ModelGTSample:
        ...

    def collate_fn(self, batch: Sequence[ModelGTSample]) -> ModelGTBatch:
        ...

T4Dataset implements it for T4 databases, and NuScenesDataset extends it for nuScenes. The annotations of each task come from a task dataset deriving BaseDatasetTask, such as Detection3DTask and T4Segmentation3DTask. A dataset config lists its task datasets under dataset_tasks, so one dataset class serves detection, segmentation or both. The dataset config leaves its transforms required (???), and every model sets the pipeline of each split in its task config.

get_data_sample() returns the sensor file paths, calibration and annotations of a record. Point clouds and images are loaded by the transforms. The collate function stacks the samples into a typed ModelGTBatch.

Transforms¶

Transforms are composable data augmentations applied per sample on CPU. Each transform takes a ModelGTSample and returns the updated sample.

class BaseTransform(ABC):
    _required_keys: Sequence[str] = ()

    def __call__(self, model_gt_sample: ModelGTSample) -> ModelGTSample:
        self._validate_required_keys(model_gt_sample)
        if not self._should_apply():
            return self.on_skip(model_gt_sample)
        return self.transform(model_gt_sample)

    @abstractmethod
    def transform(self, model_gt_sample: ModelGTSample) -> ModelGTSample:
        ...

class TransformsCompose:
    def __init__(self, pipeline: Sequence[BaseTransform]):
        self.pipeline = pipeline

    def __call__(self, model_gt_sample: ModelGTSample) -> ModelGTSample:
        for transform in self.pipeline:
            model_gt_sample = transform(model_gt_sample)
        return model_gt_sample

Every dataset config carries its own transforms, so each split has its own pipeline. The dataset applies it in __getitem__().

Public transform targets should reference the concrete implementation module, for example autoware_ml.transforms.point_cloud.loading.LoadPointsFromFile or autoware_ml.transforms.camera_lidar.geometry.GlobalBEVRandomFlip. Avoid package-level re-export layers in __init__.py. Imports and Hydra _target_ paths should point at the implementation module directly.

Runtime Data Preprocessing¶

Runtime preprocessing is a DataPreprocessor the model owns. The task config declares it at the top level as data_preprocessor, the builders instantiate it and pass it to the model's constructor. It runs on the target device after Lightning moves the batch over, and before the model's forward().

class DataPreprocessor:
    def __init__(self, preprocessor_modules: Sequence[DataPreprocessorModule]) -> None:
        self.preprocessor_modules = tuple(preprocessor_modules)

    def __call__(self, batch: ModelGTBatch, *, is_training: bool) -> ModelBatchInputs:
        batch_inputs = ModelBatchInputs.from_gt_batch(batch)
        for module in self.preprocessor_modules:
            batch_inputs = module(batch_inputs, is_training=is_training)
        return batch_inputs

Every module implements DataPreprocessorModule: it takes a ModelBatchInputs and the keyword is_training, and returns the inputs with the features it computes added or the data it changes replaced. The model inputs are a typed ModelBatchInputs. It holds the collated ModelGTBatch and one field for each kind of preprocessed feature, such as voxels_data, range_view_data and image_data.

ModuleBaseModel.on_after_batch_transfer() applies the preprocessor. Output-side shaping (e.g., logits -> probabilities, voxel-to-point scatter) lives inside the model, not in a framework pipeline: each model handles it in its own forward(), compute_metrics(), build_eval_output() and decode_outputs(). Keeping this logic in the model class avoids invisible load-bearing dependencies between config composition and metric correctness.

Model¶

All supported models inherit from ModuleBaseModel (extending LightningModule), which provides a standard interface and a set of override hooks for task-specific behavior:

class ModuleBaseModel(MetricEvalMixin, L.LightningModule, ABC):
    def __init__(
        self,
        log_dict_configs: LogDictConfigs,
        data_preprocessor: DataPreprocessor,
        optimizer: Callable[..., Optimizer] | None = None,
        scheduler: Callable[[Optimizer], LRScheduler] | None = None,
        optimizer_group_overrides: Mapping[str, Mapping[str, Any]] | None = None,
        scheduler_config: Mapping[str, Any] | None = None,
        metrics: Sequence[MetricSuite] | None = None,
    ):
        ...

    def forward(self, multi_task_batch_inputs: ModelBatchInputs) -> ModelOutputs:
        ...  # abstract

    def compute_metrics(
        self, multi_task_batch_inputs: ModelBatchInputs, multi_task_outputs: ModelOutputs
    ) -> Mapping[str, torch.Tensor]:
        ...  # abstract, returns at least "loss"

    def build_eval_output(self, batch: ModelBatchInputs, outputs: ModelOutputs) -> dict[str, Any]:
        ...  # the flat dict the metric suites read, empty by default

    def decode_outputs(
        self, multi_task_batch_inputs: ModelBatchInputs, outputs: ModelOutputs
    ) -> ModelPredictions:
        ...  # abstract, used by the predict step

    def get_log_batch_size(self, multi_task_batch_inputs: ModelBatchInputs) -> int | None:
        ...

    def build_optimizer_groups(self) -> Mapping[str, Sequence[torch.nn.Parameter]]:
        ...

    def build_export_spec(self, multi_task_batch_inputs: ModelBatchInputs) -> ExportSpec:
        ...  # required by deploy

    def build_export_specs(self, multi_task_batch_inputs: ModelBatchInputs) -> dict[str, ExportSpec]:
        ...  # one spec per deployable module, "end_to_end" by default

    def configure_optimizers(self) -> Optimizer | dict[str, Any]:
        ...

The base class handles:

• Unified step logic - All models share the same training, validation, test, and predict execution path
• Runtime data preprocessing - Applies the DataPreprocessor of the model after batch transfer
• Metric logging - Logs every entry of compute_metrics() under the stage prefix per epoch, additionally per training step when log_dict_configs asks for it, and feeds build_eval_output() into the attached metric suites at validation and test
• Predict step - Runs forward and formats predictions via decode_outputs()
• Export contract - Every model builds its own export specification from the model inputs
• Example input - Collates the first training sample as the example input so the model summary reports FLOPs

Models can have any internal architecture. forward() takes the typed ModelBatchInputs and returns a typed ModelOutputs with the head outputs of every task the model predicts. decode_outputs() decodes them into a typed ModelPredictions. Both hold one optional field per task, so a single task model fills one field and a multi task model fills several. The deployment export cannot trace those dataclasses, so build_export_spec() flattens the batch into the tensor arguments of the exported graph. Specialized models can override hooks such as decode_outputs(), build_eval_output(), get_log_batch_size(), build_optimizer_groups() or build_export_spec() without leaving the shared framework contract.

Deployment Pipeline¶

The deployment pipeline exports trained models to production-ready formats:

flowchart LR
    subgraph ONNXExport [ONNX Export]
        Checkpoint[Checkpoint] --> Load[Load Weights]
        Load --> Model[Model Eval Mode]
        Model --> Trace[Trace with Sample]
        Trace --> ONNX[ONNX File]
    end

    subgraph TensorRTBuild [TensorRT Build]
        ONNX --> Parse[Parse ONNX]
        Parse --> Optimize[Build Engine]
        Optimize --> EngineFile[Engine File]
    end

    style Checkpoint fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Load fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Model fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Trace fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style ONNX fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Parse fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style Optimize fill:#a5d6a7,opacity:0.2,stroke:#05bc23
    style EngineFile fill:#a5d6a7,opacity:0.2,stroke:#05bc23

The deployment process:

1. Load checkpoint - Instantiates model from config and loads weights from checkpoint
2. Get input sample - Uses the predict dataloader to obtain a preprocessed sample for deployment
3. Resolve export spec - Builds the effective export module and example inputs through the model's build_export_spec() contract
4. Export to ONNX - Traces the resolved export module, supporting dynamic shapes for variable input sizes
5. Build TensorRT engine - Optimizes the ONNX model for inference on NVIDIA GPUs with configurable optimization profile

Configuration is done through the deploy section in task configs.

Extending the Framework¶

See Adding Models for a detailed guide.