可插拔训练加速trick-Scaling PyTorch Model Training With Minimal Code Changes

依赖:

shell 复制代码
pip install lightning

插拔改动:

python 复制代码
from lightning.fabric import Fabric

#...

# 实例化
fabric = Fabric(accelerator='cuda')  
# 混精度用这个,加速明显
#fabric = Fabric(accelerator="cuda", precision="bf16-mixed")
fabric.launch()

#...

# 插拔接入
model, optimizer = fabric.setup(model, optimizer) 
train_dataloader = fabric.setup_dataloaders(train_dataloader)

#...

def train(num_epochs, model, optimizer, train_loader, val_loader, fabric):
    for epoch in range(num_epochs):
        train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10).to(fabric.device)

        model.train()
        for batch_idx, (features, targets) in enumerate(train_loader):
            model.train()  
            logits = model(features)
            loss = F.cross_entropy(logits, targets)
            optimizer.zero_grad()
            fabric.backward(loss)  # 插拔接入,原反向传播:loss.backward()
            optimizer.step()
            #...

参考文献

CVPR 2023 Talk:Scaling PyTorch Model Training With Minimal Code Changes

相关推荐
新智元11 分钟前
奥特曼:再也不和小扎说话!OpenAI 偷袭小扎马斯克,反手挖 4 核心员工
人工智能·openai
新智元16 分钟前
CS 专业爆冷,失业率达艺术史 2 倍!年入千万只需 5 年,大学却在禁 Cursor
人工智能·openai
代码能跑就行管它可读性21 分钟前
【论文复现】利用生成式AI进行选股和分配权重
人工智能·chatgpt
阿里云大数据AI技术30 分钟前
ODPS 15周年开发者活动|征文+动手实践双赛道开启,参与活动赢定制好礼!
大数据·人工智能·云计算
一颗小树x30 分钟前
【机器人】复现 Aether 世界模型 | 几何感知统一 ICCV 2025
人工智能·机器人·世界模型·aether
Black_Rock_br33 分钟前
语音交互新纪元:Hugging Face LeRobot如何让机器人真正“懂你”
人工智能·计算机视觉·机器人
1900_35 分钟前
【论文解读】Referring Camouflaged Object Detection
人工智能·目标检测·计算机视觉
RobotLife42 分钟前
nng库使用
人工智能
机器之心1 小时前
「世界模型」也被泼冷水了?邢波等人揭开五大「硬伤」,提出新范式
人工智能
甲丁1 小时前
国内 Claude Code 接入指南(免费获得国内代理$100额度)
人工智能