Composer 训练优化 — MosaicML 出品的训练速度和成本优化库¶

# 安装 Composer
pip install mosaicml                       # 安装 Composer（包名是 mosaicml）
pip install composer                       # 或者用这个包名（同一个库）

# 含所有可选依赖
pip install 'mosaicml[all]'

# 验证安装
python -c "import composer; print(composer.__version__)"

import composer
from composer import Trainer                            # Composer 的核心训练器
from composer.models import ComposerModel               # 模型包装基类
from composer.algorithms import (                       # 各种训练加速算法
    LabelSmoothing,                                     # 标签平滑（防止过拟合）
    GradientClipping,                                   # 梯度裁剪（防止梯度爆炸）
    BlurPool,                                           # 模糊池化（提升泛化）
    ProgressiveResizing,                                # 渐进式分辨率（CV 专用）
)

# 包装 PyTorch 模型（需要继承 ComposerModel）
class MyModel(ComposerModel):
    def __init__(self):
        super().__init__()
        self.model = ...                                # 你的 PyTorch 模型

    def forward(self, batch):                           # 前向传播
        x, y = batch
        return self.model(x)

    def loss(self, outputs, batch):                     # 计算 loss
        x, y = batch
        return torch.nn.functional.cross_entropy(outputs, y)

    def get_metrics(self, is_train):                    # 评估指标
        return {"accuracy": Accuracy()}

# 创建 Trainer（指定各种优化算法）
trainer = Trainer(
    model=MyModel(),
    train_dataloader=train_loader,
    eval_dataloader=val_loader,
    max_duration="10ep",                               # 训练 10 个 epoch
    algorithms=[                                       # 要使用的优化算法列表
        LabelSmoothing(smoothing=0.1),                 # 标签平滑 0.1
        GradientClipping(clipping_type="norm",         # 梯度裁剪类型
                         clipping_threshold=1.0),      # 裁剪阈值
    ],
    device="gpu",                                      # 使用 GPU
    precision="amp_bf16",                              # 自动混合精度（bf16）
)

trainer.fit()

from composer import Trainer
from composer.callbacks import (                        # 训练回调（监控和控制训练过程）
    LRMonitor,                                          # 记录学习率变化
    MemoryMonitor,                                      # 监控显存使用
    SpeedMonitor,                                       # 监控训练速度
    CheckpointSaver,                                    # 自动保存检查点
)
from composer.loggers import WandBLogger                # W&B 实验追踪

trainer = Trainer(
    model=model,
    train_dataloader=train_loader,
    eval_dataloader=val_loader,

    # 训练时长（支持多种格式）
    max_duration="10ep",                               # 10 个 epoch
    # max_duration="10000ba",                          # 10000 个 batch
    # max_duration="2hr",                              # 2 小时

    # 精度设置
    precision="amp_bf16",                              # amp_fp16/amp_bf16/fp32

    # 混合精度梯度缩放
    grad_clip_norm=1.0,                                # 梯度裁剪范数
    grad_accum="auto",                                 # 自动梯度累积（适配显存）

    # 回调和日志
    callbacks=[
        LRMonitor(),                                    # 记录学习率
        SpeedMonitor(window_size=50),                   # 每 50 步计算一次速度
        CheckpointSaver(                                # 检查点保存策略
            folder="./checkpoints",
            save_interval="500ba",                      # 每 500 batch 保存
            num_checkpoints_to_keep=3,                  # 只保留最近 3 个
        ),
    ],
    loggers=[
        WandBLogger(project="my_training"),             # W&B 追踪
    ],

    # 优化器（如果不指定，使用 DecoupledSGDW）
    optimizers=torch.optim.AdamW(model.parameters(), lr=1e-4),
    schedulers=composer.optim.LinearWithWarmupScheduler(
        t_warmup="100ba",                              # 前 100 个 batch 预热
    ),
)

# 用 Composer 训练图像分类模型（ResNet + 多种加速算法）

import composer
from composer import Trainer
from composer.models import ComposerResNet              # 内置的 ResNet 包装
from composer.datasets import get_cifar10_dataloader    # 内置数据集加载器
from composer.algorithms import (
    ProgressiveResizing,                               # 从小分辨率逐渐增大
    MixUp,                                             # MixUp 数据增强
    CutMix,                                            # CutMix 数据增强
    ColOut,                                            # 列级数据增广
    LabelSmoothing,
)

model = ComposerResNet(model_name="resnet50", num_classes=10)

train_loader = get_cifar10_dataloader(
    datadir="/data/cifar10",
    is_train=True,
    batch_size=256,
    num_workers=8,
)

trainer = Trainer(
    model=model,
    train_dataloader=train_loader,
    max_duration="90ep",                               # CIFAR-10 标准 90 轮
    algorithms=[
        ProgressiveResizing(initial_scale=0.5),        # 从 50% 分辨率开始
        MixUp(alpha=0.2),                              # MixUp α=0.2
        LabelSmoothing(smoothing=0.1),
    ],
    device="gpu",
    precision="amp_bf16",
)
trainer.fit()

# 导出模型
trainer.save_checkpoint("./final_checkpoint.pt")       # 保存最终检查点