可插拔训练加速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

相关推荐
星期天要睡觉17 分钟前
计算机视觉(opencv)实战二十一——基于 SIFT 和 FLANN 的指纹图像匹配与认证
人工智能·opencv·计算机视觉
victory043120 分钟前
wav2vec微调进行疾病语音分类任务
人工智能·分类·数据挖掘
semantist@语校28 分钟前
第二十篇|SAMU教育学院的教育数据剖析:制度阈值、能力矩阵与升学网络
大数据·数据库·人工智能·百度·语言模型·矩阵·prompt
IT_陈寒1 小时前
React 性能优化必杀技:这5个Hook组合让你的应用提速50%!
前端·人工智能·后端
剪一朵云爱着1 小时前
一文入门:机器学习
人工智能·机器学习
hi0_61 小时前
机器学习实战(一): 什么是机器学习
人工智能·机器学习·机器人·机器学习实战
ChinaRainbowSea1 小时前
9. LangChain4j + 整合 Spring Boot
java·人工智能·spring boot·后端·spring·langchain·ai编程
有Li2 小时前
基于联邦学习与神经架构搜索的可泛化重建:用于加速磁共振成像|文献速递-最新医学人工智能文献
论文阅读·人工智能·文献·医学生
桃花键神2 小时前
从传统到智能:3D 建模流程的演进与 AI 趋势 —— 以 Blender 为例
人工智能·3d·blender
星期天要睡觉2 小时前
计算机视觉(opencv)实战十七——图像直方图均衡化
人工智能·opencv·计算机视觉