repository·main·Indexed 23 days ago
https://github.com/open-mmlab/mmengineA foundational library for the OpenMMLab 2.0 ecosystem. It provides core components such as the BaseDataset class for managing metadata and data pipelines, as well as support for Docker-based production and development environments. The library includes examples for Llama2 training with FSDP fine-tuning, image segmentation, text classification, and text translation.
The mmengine.logging module provides tools for managing logs, message hubs, and history buffers during execution. The core components are:
MMLogger: The primary interface for logging messages.MessageHub: A component used to collect and manage messages (often used for tracking metrics or status updates).HistoryBuffer: A buffer used to store historical log information.print_log: A utility function for printing log messages.The mmengine.structures module provides specialized data structures designed to handle common data types in computer vision and machine learning tasks. These structures are optimized for representing elements like instances, labels, and pixel-level data, facilitating efficient data management during training and inference.
Key data structures include:
BaseDataElement: The base class for all data elements in this module.InstanceData: Used for representing object instances (e.g., bounding boxes, scores).LabelData: Used for representing ground truth or predicted labels.PixelData: Used for representing pixel-wise information (e.g., segmentation masks).The mmengine.hooks module provides a set of hook classes designed to inject custom logic into the training or evaluation lifecycle of an OpenMMLab-based runner. Hooks allow you to perform tasks such as logging, checkpointing, parameter scheduling, and visualization at specific stages (e.g., before/after an iteration or epoch) without modifying the core engine logic.
Key categories of hooks include:
CheckpointHook (saving models), EarlyStoppingHook (stopping training based on metrics), and ParamSchedulerHook (adjusting parameters).LoggerHook (recording metrics), IterTimerHook (measuring iteration time), and RuntimeInfoHook (tracking system status).EMAHook (Exponential Moving Average), SyncBuffersHook (synchronizing buffers in distributed training), and EmptyCacheHook (managing GPU memory).NaiveVisualizationHook (visualizing data/results), ProfilerHook, and NPUProfilerHook (performance profiling).The MMEngine logging system is designed to manage and visualize various logs generated during a Runner's execution (e.g., dataset info, model initialization, learning rate, loss).
Key components include:
MessageHub (mmengine.logging.MessageHub): Maintains the current training state.HistoryBuffer (mmengine.logging.HistoryBuffer): Manages and encapsulates historical log data (trajectories).LogProcessor (mmengine.runner.LogProcessor): Formats data from the MessageHub.MMLogger (mmengine.logging.MMLogger): The logger interface.LoggerHook (mmengine.hooks.LoggerHook): Displays formatted logs to various visualization backends.Note: Most users do not need to interact with the internal data flow directly; instead, they can configure the LogProcessor via configuration files to choose preferred log statistics (e.g., average loss over a window vs. smoothed loss).
mmengine.dataset module provides core abstractions for handling datasets in machine learning workflows. It includes base classes for defining custom datasets, wrappers for manipulating existing datasets (such as concatenation or balancing), samplers for controlling data access patterns, and utilities for data loading and collation.The mmengine.analysis module provides tools to analyze the computational complexity of models, including parameter counts, FLOPs (Floating Point Operations), and activation counts.
Key components include:
ActivationAnalyzer and FlopAnalyzer for detailed analysis.The mmengine.dataset module provides core abstractions for handling data in machine learning workflows.
BaseDataset: The foundation for all dataset implementations.Compose: Used to compose multiple data transformations or operations into a single pipeline.Wrappers allow you to modify the behavior of an existing dataset without changing its underlying implementation:
ClassBalancedDataset: Adjusts the dataset to handle class imbalance.ConcatDataset: Concatenates multiple datasets into one.RepeatDataset: Repeats a dataset a specified number of times.Samplers control how data indices are selected during training or evaluation:
DefaultSampler: The standard sampler for iterating through data.InfiniteSampler: A sampler that provides an infinite stream of data indices, useful for certain training loops.mmengine.model module provides core abstractions for building, wrapping, and initializing neural network models. It includes base classes for modules, models, EMA (Exponential Moving Average) strategies, and distributed data parallel wrappers. It also provides a comprehensive suite of weight initialization methods and utility functions for model-related operations.The mmengine.registry module provides the core mechanism for component registration and management in MMEngine. It allows you to register classes, functions, or other objects into a central registry, which can then be instantiated using configuration objects.
Key components include:
Registry: The main class used to create and manage registries for specific types of modules (e.g., models, datasets, optimizers).DefaultScope: Manages the default scope for resolving module names, allowing for hierarchical organization and easier instantiation via configuration.Common utility functions for building objects from configurations include:
build_from_cfg: General purpose utility to build an object from a configuration.build_model_from_cfg: Specifically for building models.build_runner_from_cfg: Specifically for building runners.build_scheduler_from_cfg: Specifically for building schedulers.The mmengine.optim module provides high-level wrappers for PyTorch optimizers and a wide range of learning rate and momentum schedulers.
mmengine.optim includes several wrapper classes to manage optimizer state and behavior, such as:
OptimWrapper: The base wrapper for optimizers.OptimWrapperDict: For managing multiple optimizers.AmpOptimWrapper: For Automatic Mixed Precision (AMP) training.ApexOptimWrapper: For NVIDIA Apex-based mixed precision.ZeroRedundancyOptimizer: For ZeRO-style redundancy reduction.build_optim_wrapper: A utility function to construct these wrappers.mmengine.optim provides extensive support for parameter and momentum scheduling. Schedulers are categorized by their decay pattern:
CosineAnnealingLR, ExponentialLR, LinearLR, MultiStepLR, OneCycleLR, PolyLR, StepLR, ReduceOnPlateauLR, etc.CosineAnnealingMomentum, ExponentialMomentum, LinearMomentum, MultiStepMomentum, etc.CosineAnnealingParamScheduler, ExponentialParamScheduler, LinearParamScheduler, etc.ConstantLR, ConstantMomentum, ConstantParamScheduler.mmengine.runner module provides the core execution logic for training, validation, and testing workflows. It includes the Runner classes for managing the overall lifecycle, Loop abstractions for defining specific execution steps (like training or validation), and utilities for checkpoint management and Automatic Mixed Precision (AMP).The mmengine.dist module provides a suite of collective communication primitives for distributed training and multi-GPU environments. These functions wrap standard distributed backend operations (like NCCL or Gloo) to facilitate data synchronization across different processes.
gather or gather_object to collect tensors or Python objects from all processes to a single process.all_gather or all_gather_object to collect data from all processes and distribute the complete set back to all processes.all_reduce, all_reduce_dict, or all_reduce_params to perform mathematical reductions (like sum or mean) across all processes.broadcast or broadcast_object_list to send data from a source process to all other processes.sync_random_seed to ensure all processes use the same random seed for reproducibility, and collect_results (with CPU/GPU variants) to aggregate evaluation results.