【框架】简化多卡训练——huggingface accelerate使用方法介绍

最新推荐文章于 2025-03-26 14:10:09 发布

Big Bill Fighter

最新推荐文章于 2025-03-26 14:10:09 发布

阅读量1.3w

点赞数 19

分类专栏： DeepLearning 文章标签： linux 深度学习 pytorch deep learning 人工智能 python 服务器

本文链接：https://blog.csdn.net/weixin_43441566/article/details/139061337

版权

HuggingFace 的 accelerate 库可以实现只需要修改几行代码就可以实现ddp训练，且支持混合精度训练和TPU训练。(甚至支持deepspeed。)
accelerate支持的训练方式为CPU/单GPU (TPU)/多GPU(TPU) DDP模式/fp32/fp16等。

安装

pip install accelerate

使用

使用accelerate进行单卡或者多卡训练的代码是相同的，不过在单卡训练的时候可以不使用gather_for_metrics()函数聚合信息。这里为了保持代码的不变性，仍然保留gather_for_metrics()。以下为使用accelerate运行在MNIST数据集上面运行手写数字识别的样例代码main.py。

import datetime
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

#======================================================================
# import accelerate
from accelerate import Accelerator
from accelerate.utils import set_seed
#======================================================================


class BasicNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.act = F.relu

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)

最低0.47元/天解锁文章