graph LR
RunManager["RunManager"]
DistributedRunManager["DistributedRunManager"]
Progressive_Shrinking_Module["Progressive Shrinking Module"]
ImageNet_Evaluation_Helper_Module["ImageNet Evaluation Helper Module"]
Network["Network"]
Optimizer["Optimizer"]
Distributed_Backend["Distributed Backend"]
RunManager -- "interacts with" --> Network
RunManager -- "interacts with" --> Optimizer
DistributedRunManager -- "inherits from" --> RunManager
DistributedRunManager -- "interacts with" --> Distributed_Backend
Progressive_Shrinking_Module -- "delegates to" --> RunManager
Progressive_Shrinking_Module -- "delegates to" --> DistributedRunManager
Progressive_Shrinking_Module -- "interacts with" --> Network
ImageNet_Evaluation_Helper_Module -- "uses" --> RunManager
ImageNet_Evaluation_Helper_Module -- "uses" --> DistributedRunManager
ImageNet_Evaluation_Helper_Module -- "interacts with" --> Network
The OFA (Once-For-All) project's core training and evaluation subsystem is orchestrated by the RunManager and its distributed counterpart, DistributedRunManager. These managers oversee the training lifecycle, interacting with the Network component (represented by elastic neural networks like OFAMobileNetV3) and an abstract Optimizer to update model weights. The Progressive Shrinking Module leverages these run managers to implement the multi-stage training strategy for the supernet, dynamically adjusting network parameters. For evaluation, the ImageNet Evaluation Helper Module provides specialized utilities, utilizing the run managers and interacting directly with the Network to assess subnet performance on the ImageNet dataset. The Distributed Backend facilitates parallel processing, primarily through the DistributedRunManager, ensuring efficient and synchronized operations in distributed training environments. This architecture enables flexible and scalable neural architecture search and deployment.
The foundational orchestrator for standard (non-distributed) training and evaluation loops. It manages the core training process, validation, model checkpointing, and progress logging.
Related Classes/Methods:
Extends RunManager to coordinate training and evaluation across multiple devices or nodes, ensuring synchronized model updates and efficient data distribution in a distributed environment.
Related Classes/Methods:
Implements the multi-stage progressive shrinking strategy, which is crucial for training the supernet. It orchestrates a sequence of training phases (e.g., elastic depth, width, kernel size) to gradually refine the supernet for elasticity.
Related Classes/Methods:
Provides specialized utilities for evaluating OFA subnets specifically on the ImageNet dataset. This includes pre-evaluation steps like BatchNorm calibration and specific subnet configuration for accurate performance measurement.
Related Classes/Methods:
Represents the neural network model, specifically the elastic OFA networks that can be configured with varying architectural parameters.
Related Classes/Methods:
Manages the optimization process for model weights during training. While not a single dedicated class within the OFA project, its functionality is integrated and managed by the RunManager and DistributedRunManager components, which typically leverage standard deep learning optimizers (e.g., from PyTorch).
Related Classes/Methods: None
Handles communication and synchronization in distributed environments, enabling training across multiple devices or nodes. This component is primarily interacted with and abstracted by the DistributedRunManager.
Related Classes/Methods: