graph LR
Training_Evaluation_Orchestrator["Training & Evaluation Orchestrator"]
Optimizer_Management["Optimizer Management"]
Loss_Function_Library["Loss Function Library"]
Metric_Aggregation_Utility["Metric Aggregation Utility"]
Training_Evaluation_Orchestrator -- "configures optimizers through" --> Optimizer_Management
Training_Evaluation_Orchestrator -- "utilizes" --> Loss_Function_Library
Training_Evaluation_Orchestrator -- "updates" --> Metric_Aggregation_Utility
Loss_Function_Library -- "provides calculations to" --> Training_Evaluation_Orchestrator
Metric_Aggregation_Utility -- "provides metrics to" --> Training_Evaluation_Orchestrator
The Training & Optimization Engine subsystem is central to CompressAI, providing the necessary framework for training and optimizing compression models. It adheres to the "ML Toolkit/Library with an Evaluation Platform" architectural patterns by modularizing core training logic, optimizer management, and loss function definitions.
This component manages the entire training and evaluation lifecycle. It orchestrates the execution of training and testing epochs, handles model checkpointing, and coordinates the overall flow of the training process. It acts as the primary control unit for model training experiments.
Related Classes/Methods:
examples.train:mainexamples.train_pointcloud:mainexamples.train_video:mainexamples.train:train_one_epochexamples.train_pointcloud:train_one_epochexamples.train_video:train_one_epochexamples.train:test_epochexamples.train_pointcloud:test_epochexamples.train_video:test_epoch
This component is responsible for the instantiation, configuration, and management of various PyTorch optimizers. It provides standardized methods for creating optimizers for both the main network and any auxiliary components, ensuring consistency and reusability across different training scripts.
Related Classes/Methods:
examples.train:configure_optimizersexamples.train_pointcloud:configure_optimizersexamples.train_video:configure_optimizerscompressai.optimizers.net:net_optimizercompressai.optimizers.net:make_optimizercompressai.optimizers.net_aux:net_aux_optimizercompressai.optimizers.net_aux:make_optimizer
This component defines and implements a collection of loss functions specifically tailored for different data compression models (e.g., point cloud, video). It encapsulates the logic for calculating rate-distortion, reconstruction, and auxiliary losses, providing a flexible and extensible set of objectives for model training.
Related Classes/Methods:
compressai.losses.pointcloud.chamfercompressai.losses.pointcloud.hrtzxf2022examples.train_video:RateDistortionLossexamples.train_video:compute_aux_loss
This utility component is designed to accumulate and average numerical metrics over a specified period, typically an epoch. It provides a standardized way to track and summarize model performance statistics, such as loss values, PSNR, or other relevant evaluation metrics.
Related Classes/Methods: