PyTorch——利用Accelerate轻松控制多个CPU/GPU/TPU加速计算

news2025/2/25 18:01:16

PyTorch——利用Accelerate轻松控制多个CPU/GPU/TPU加速计算

    • 前言
    • 官方示例
    • 控制多个CPU/GPU/TPU
      • 简单说一下
      • 设备环境
      • 导包
      • 加载数据 FashionMNIST
      • 创建一个简单的CNN模型
      • 训练函数-只包含训练
      • 训练函数-包含训练和验证
      • 训练
    • 参考链接

前言

  • CPU?GPU?TPU?
    • 计算设备太多,很混乱?
    • 切换环境,代码大量改来改去?
    • 不懂怎么调用多个CPU/GPU/TPU?或者想轻松调用?
  • OK!OK!OK!
    • 来自HuggingFace的Accelerate库帮你轻松解决这些问题,只需几行代码改动既可以快速完成计算设备的自动调整。
      huggingface
  • 相关地址
    • 官方文档:https://huggingface.co/docs/accelerate/index
    • GitHub:https://github.com/huggingface/accelerate
    • 安装(推荐用>=0.14的版本) $ pip install accelerate
  • 下面就来说说怎么用
    • 你也可以直接看我在Kaggle上做好的完整的Notebook示例

官方示例

  • 先大致看个样
  • 移除掉以前.to(device)部分的代码,引入Acceleratormodel、optimizer、data、loss.backward()做下处理即可
import torch
import torch.nn.functional as F
from datasets import load_dataset
from accelerate import Accelerator

# device = 'cpu'
accelerator = Accelerator()

# model = torch.nn.Transformer().to(device)
model = torch.nn.Transformer()
optimizer = torch.optim.Adam(model.parameters())

dataset = load_dataset('my_dataset')
data = torch.utils.data.DataLoader(dataset, shuffle=True)

model, optimizer, data = accelerator.prepare(model, optimizer, data)

model.train()
for epoch in range(10):
    for source, targets in data:
        # source = source.to(device)
        # targets = targets.to(device)

        optimizer.zero_grad()

        output = model(source)
        loss = F.cross_entropy(output, targets)

        # loss.backward()
        accelerator.backward(loss)

        optimizer.step()

控制多个CPU/GPU/TPU

简单说一下

  • 对于单个设备,像前面那个简单示例改下代码即可
  • 多个设备的情况下,有一点特殊的要处理,下面做个完整的PyTorch训练示例
    • 你可以拿这个和我之前发的示例做个对比 CNN图像分类-FashionMNIST
    • 也可以直接看我在Kaggle上做好的完整的Notebook示例

设备环境

  • 看看当前的显卡设备(2颗Tesla T4),命令 $ nvidia-smi
Thu Apr 27 10:53:26 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   43C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla T4            Off  | 00000000:00:05.0 Off |                    0 |
| N/A   41C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
  • 安装或更新Accelerate,命令 $ !pip install --upgrade accelerate

导包

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor, Compose
import torchvision.datasets as datasets
from accelerate import Accelerator
from accelerate import notebook_launcher

加载数据 FashionMNIST

train_data = datasets.FashionMNIST(
    root="./data",
    train=True,
    download=True,
    transform=Compose([ToTensor()])
)

test_data = datasets.FashionMNIST(
    root="./data",
    train=False,
    download=True,
    transform=Compose([ToTensor()])
)

print(train_data.data.shape)
print(test_data.data.shape)

创建一个简单的CNN模型

class CNNModel(nn.Module):
    
  def __init__(self):
    super(CNNModel, self).__init__()
    self.module1 = nn.Sequential(
        nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )  
    self.module2 = nn.Sequential(
        nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )
    self.flatten = nn.Flatten()
    self.linear1 = nn.Linear(7 * 7 * 64, 64)
    self.linear2 = nn.Linear(64, 10)
    self.relu = nn.ReLU()
    
  def forward(self, x):
    out = self.module1(x)
    out = self.module2(out)
    out = self.flatten(out)
    out = self.linear1(out)
    out = self.relu(out)
    out = self.linear2(out)
    return out

训练函数-只包含训练

  • 注意看accelerator相关代码
  • 若要实现多设备控制训练,for epoch in range(epoch_num):中末尾处的代码必不可少
def training_function():
    # 参数配置
    epoch_num = 4
    batch_size = 64
    learning_rate = 0.005

    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    
    # 数据
    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)

    # 模型/损失函数/优化器
    # model = CNNModel().to(device)
    model = CNNModel()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    accelerator = Accelerator()
    model, optimizer, train_loader, val_loader = accelerator.prepare(model, optimizer, train_loader, val_loader)

    # 开始训练
    for epoch in range(epoch_num):
        # 训练
        model.train()
        for i, (X_train, y_train) in enumerate(train_loader):
            # X_train = X_train.to(device)
            # y_train = y_train.to(device)
            out = model(X_train)
            loss = criterion(out, y_train)

            optimizer.zero_grad()
            # loss.backward()
            accelerator.backward(loss)
            optimizer.step()

            if (i + 1) % 100 == 0:
                print(f"{accelerator.device} Train... [epoch {epoch + 1}/{epoch_num}, step {i + 1}/{len(train_loader)}]\t[loss {loss.item()}]")
        
        # 等待每个GPU上的模型执行完当前的epoch,并进行合并同步
        accelerator.wait_for_everyone() 
        model = accelerator.unwrap_model(model)
        # 现在所有GPU上都一样了,可以保存model
        accelerator.save(model, "model.pth") 

训练函数-包含训练和验证

  • 相比前面的代码,多了“验证”相关的代码
  • 验证时,因为使用多个设备进行训练,所以会比较特殊,会涉及到多个设备的验证结果合并的问题
def training_function():
    # 参数配置
    epoch_num = 4
    batch_size = 64
    learning_rate = 0.005

    # 数据
    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)

    # 模型/损失函数/优化器
    model = CNNModel()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    accelerator = Accelerator()
    model, optimizer, train_loader, val_loader = accelerator.prepare(model, optimizer, train_loader, val_loader)

    # 开始训练
    for epoch in range(epoch_num):
        # 训练
        model.train()
        for i, (X_train, y_train) in enumerate(train_loader):
            out = model(X_train)
            loss = criterion(out, y_train)

            optimizer.zero_grad()
            accelerator.backward(loss)
            optimizer.step()

            if (i + 1) % 100 == 0:
                print(f"{accelerator.device} Train... [epoch {epoch + 1}/{epoch_num}, step {i + 1}/{len(train_loader)}]\t[loss {loss.item()}]")
                
        # 验证
        model.eval()
        correct, total = 0, 0
        for X_val, y_val in val_loader:
            with torch.no_grad():
                output = model(X_val)
            _, pred = torch.max(output, 1)
            
            # 合并每个GPU的验证数据
            pred, y_val = accelerator.gather_for_metrics((
                pred, y_val
            ))
            
            total += y_val.size(0)
            correct += (pred == y_val).sum()

        # 用main process打印accuracy
        accelerator.print(f'epoch {epoch + 1}/{epoch_num}, accuracy = {100 * (correct.item() / total):.2f}')
        
        # 等待每个GPU上的模型执行完当前的epoch,并进行合并同步
        accelerator.wait_for_everyone() 
        model = accelerator.unwrap_model(model)
        # 现在所有GPU上都一样了,可以保存model
        accelerator.save(model, "model.pth") 

训练

  • 如果你在本地训练的话,直接调用前面定义的函数training_function即可
training_function()
  • 如果你在Kaggle/Colab上面,则需要利用notebook_launcher进行训练
# num_processes=2 指定使用2个GPU,因为当前我申请了2颗 Nvidia T4
notebook_launcher(training_function, num_processes=2)
  • 下面是2个GPU训练时的控制台输出样例
Launching training on 2 GPUs.
cuda:0 Train... [epoch 1/4, step 100/469]	[loss 0.43843933939933777]
cuda:1 Train... [epoch 1/4, step 100/469]	[loss 0.5267877578735352]
cuda:0 Train... [epoch 1/4, step 200/469]	[loss 0.39918822050094604]cuda:1 Train... [epoch 1/4, step 200/469]	[loss 0.2748252749443054]

cuda:1 Train... [epoch 1/4, step 300/469]	[loss 0.54105544090271]cuda:0 Train... [epoch 1/4, step 300/469]	[loss 0.34716445207595825]

cuda:1 Train... [epoch 1/4, step 400/469]	[loss 0.2694844901561737]
cuda:0 Train... [epoch 1/4, step 400/469]	[loss 0.4343942701816559]
epoch 1/4, accuracy = 88.49
cuda:0 Train... [epoch 2/4, step 100/469]	[loss 0.19695354998111725]
cuda:1 Train... [epoch 2/4, step 100/469]	[loss 0.2911057770252228]
cuda:0 Train... [epoch 2/4, step 200/469]	[loss 0.2948791980743408]
cuda:1 Train... [epoch 2/4, step 200/469]	[loss 0.292676717042923]
cuda:0 Train... [epoch 2/4, step 300/469]	[loss 0.222089946269989]
cuda:1 Train... [epoch 2/4, step 300/469]	[loss 0.28814008831977844]
cuda:0 Train... [epoch 2/4, step 400/469]	[loss 0.3431250751018524]
cuda:1 Train... [epoch 2/4, step 400/469]	[loss 0.2546379864215851]
epoch 2/4, accuracy = 87.31
cuda:1 Train... [epoch 3/4, step 100/469]	[loss 0.24118559062480927]cuda:0 Train... [epoch 3/4, step 100/469]	[loss 0.363821804523468]

cuda:0 Train... [epoch 3/4, step 200/469]	[loss 0.36783623695373535]
cuda:1 Train... [epoch 3/4, step 200/469]	[loss 0.18346744775772095]
cuda:0 Train... [epoch 3/4, step 300/469]	[loss 0.23459288477897644]
cuda:1 Train... [epoch 3/4, step 300/469]	[loss 0.2887689769268036]
cuda:0 Train... [epoch 3/4, step 400/469]	[loss 0.3079166114330292]
cuda:1 Train... [epoch 3/4, step 400/469]	[loss 0.18255220353603363]
epoch 3/4, accuracy = 88.46
cuda:1 Train... [epoch 4/4, step 100/469]	[loss 0.27428603172302246]
cuda:0 Train... [epoch 4/4, step 100/469]	[loss 0.17705145478248596]
cuda:1 Train... [epoch 4/4, step 200/469]	[loss 0.2811894416809082]
cuda:0 Train... [epoch 4/4, step 200/469]	[loss 0.22682836651802063]
cuda:0 Train... [epoch 4/4, step 300/469]	[loss 0.2291710525751114]
cuda:1 Train... [epoch 4/4, step 300/469]	[loss 0.32024848461151123]
cuda:0 Train... [epoch 4/4, step 400/469]	[loss 0.24648766219615936]
cuda:1 Train... [epoch 4/4, step 400/469]	[loss 0.0805584192276001]
epoch 4/4, accuracy = 89.38
  • 下面是1个TPU训练时的控制台输出样例
Launching training on CPU.
xla:0 Train... [epoch 1/4, step 100/938]	[loss 0.6051161289215088]
xla:0 Train... [epoch 1/4, step 200/938]	[loss 0.27442359924316406]
xla:0 Train... [epoch 1/4, step 300/938]	[loss 0.557417631149292]
xla:0 Train... [epoch 1/4, step 400/938]	[loss 0.1840067058801651]
xla:0 Train... [epoch 1/4, step 500/938]	[loss 0.5252436399459839]
xla:0 Train... [epoch 1/4, step 600/938]	[loss 0.2718536853790283]
xla:0 Train... [epoch 1/4, step 700/938]	[loss 0.2763175368309021]
xla:0 Train... [epoch 1/4, step 800/938]	[loss 0.39897507429122925]
xla:0 Train... [epoch 1/4, step 900/938]	[loss 0.28720396757125854]
epoch = 0, accuracy = 86.36
xla:0 Train... [epoch 2/4, step 100/938]	[loss 0.24496735632419586]
xla:0 Train... [epoch 2/4, step 200/938]	[loss 0.37713131308555603]
xla:0 Train... [epoch 2/4, step 300/938]	[loss 0.3106330633163452]
xla:0 Train... [epoch 2/4, step 400/938]	[loss 0.40438592433929443]
xla:0 Train... [epoch 2/4, step 500/938]	[loss 0.38303741812705994]
xla:0 Train... [epoch 2/4, step 600/938]	[loss 0.39199298620224]
xla:0 Train... [epoch 2/4, step 700/938]	[loss 0.38932573795318604]
xla:0 Train... [epoch 2/4, step 800/938]	[loss 0.26298171281814575]
xla:0 Train... [epoch 2/4, step 900/938]	[loss 0.21517205238342285]
epoch = 1, accuracy = 90.07
xla:0 Train... [epoch 3/4, step 100/938]	[loss 0.366019606590271]
xla:0 Train... [epoch 3/4, step 200/938]	[loss 0.27360212802886963]
xla:0 Train... [epoch 3/4, step 300/938]	[loss 0.2014923095703125]
xla:0 Train... [epoch 3/4, step 400/938]	[loss 0.21998485922813416]
xla:0 Train... [epoch 3/4, step 500/938]	[loss 0.28129786252975464]
xla:0 Train... [epoch 3/4, step 600/938]	[loss 0.42534705996513367]
xla:0 Train... [epoch 3/4, step 700/938]	[loss 0.22158119082450867]
xla:0 Train... [epoch 3/4, step 800/938]	[loss 0.359947144985199]
xla:0 Train... [epoch 3/4, step 900/938]	[loss 0.3221997022628784]
epoch = 2, accuracy = 90.36
xla:0 Train... [epoch 4/4, step 100/938]	[loss 0.2814193069934845]
xla:0 Train... [epoch 4/4, step 200/938]	[loss 0.16465164721012115]
xla:0 Train... [epoch 4/4, step 300/938]	[loss 0.2897304892539978]
xla:0 Train... [epoch 4/4, step 400/938]	[loss 0.13403896987438202]
xla:0 Train... [epoch 4/4, step 500/938]	[loss 0.1135573536157608]
xla:0 Train... [epoch 4/4, step 600/938]	[loss 0.14964193105697632]
xla:0 Train... [epoch 4/4, step 700/938]	[loss 0.20239461958408356]
xla:0 Train... [epoch 4/4, step 800/938]	[loss 0.23625142872333527]
xla:0 Train... [epoch 4/4, step 900/938]	[loss 0.3418393135070801]
epoch = 3, accuracy = 90.11

参考链接

  • https://github.com/huggingface/accelerate
  • https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate
  • https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/469001.html

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!

相关文章

边学边记——数据结构☞关于对象的比较(包括对equals()方法的重写,Comparable接口,Comparator接口的介绍)

目录 一、基本类型 二、引用类型——对象的比较 1. 关于同一性的比较 2. 关于相等性的比较 三、Comparable接口 1. 介绍 2. 实现 3. 什么叫做正确的compareTo方法的重写 四、实现Comparator接口(基于比较器比较) 1. 介绍 2. 实现 3. 使用场景…

CentOS防火墙的常用快捷命令

CentOS是免费开源的Linux发行版之一,它兼容RHEL并由社区进行维护,大多数美国服务器提供对该系统支持。在使用CentOS系统时,您需要了解一些常用命令,比如开启、查看、关闭防火墙等。本文将介绍下CentOS防火墙的常用命令。 CentOS是一种面向企业级服务器环境的Linux发行版,…

D触发器仿真实验

关于D触发器的内容见专栏的单片机原理及应用,主要是时钟脉冲出现时候,会改变输出状态。 下面来做一个D触发器的仿真实验。 部件 使用74LS74:带清除和预置端功能的双上升沿D型触发器 74LS74是一款现代集成电路芯片,属于TTL&…

火爆全网,测试用例技巧-Python实现XMind测试用例转Excel测试用例

目录:导读 前言一、Python编程入门到精通二、接口自动化项目实战三、Web自动化项目实战四、App自动化项目实战五、一线大厂简历六、测试开发DevOps体系七、常用自动化测试工具八、JMeter性能测试九、总结(尾部小惊喜) 前言 自动化测试&#x…

Jenkins中Changelog插件使用

Jenkins中Changelog插件使用 通过changlog插件获取每次提交的变更信息 一、安装changelog插件 二、使用changelog 1、新建项目,在构建环境中勾选Add Changelog Information to Environment ,输入对应的格式 %3$s(datetime-at%4$s via%1$s) %2$s(type%…

聚焦智慧燃气,美格智能亮相第25届中国国际燃气、供热技术与设备展览会

4月25~27日,第25届中国国际燃气、供热技术与设备展览会在深圳会展中心(福田)盛大举办。本次展会由中国城市燃气协会主办,自1994年举办以来已历经24届,是国内燃气行业规模巨大的综合性专业展会。美格智能携4G/5G模组、N…

Typora中的数学公式(MarkDown)

Typora中的数学公式(MarkDown) 添加数学公式方式: 1.选择:段落>公式 2.快捷键:Ctrl Shift M 3.直接输入两个美元符号并回车 展示: 空格:一个斜线\ : a \ b a b a \ b a b 换行&#x…

【河南省第二届技能大赛-物联网技术】C模块Python开发讲解

文章目录 前言题目如下什么是云平台APIAPI调用的基本流程Python实例介绍依赖库介绍API调用的基本参数请求返回值requests库PyQt5库总结源码链接 前言 这是基于样题进行讲解的文章,因为正式赛题尚未公布和样题类型相似。文章结尾提供了源码和环境链接,读…

四面阿里,成功入职阿里测试开发,分享我的真实面试题

闲话少叙 直接上干货 鉴于篇幅所限,这里不放答案,有需要的朋友可以评论区自取 1. 请自我介绍一下(需简单清楚的表述自已的基本情况,在这过程中要展现出自信,对工作有激情,上进,好学) 2. 平时工作中是怎么去…

【面试题】你都必须得掌握的vue知识

大厂面试题分享 面试题库 前后端面试题库 (面试必备) 推荐:★★★★★ 地址:前端面试题库 web前端面试题库 VS java后端面试题库大全 前言 大家好,我是前端贰货道士。最近抽空整理了下我对vue2.x的理解和认知&#…

设计师常用的素材网站有哪个推荐

即时设计资源社区聚集了许多优秀的创作者,分享了大量的优质资源。 目前,社区资源数量已达到10000,包含图标、插画、原型、设计作品等多个素材类别。这些优秀的设计作品降低了设计师思维的成本,成为设计师的宝藏材料网站。 即时设…

ML@sklearn@ML流程Part2@数据划分@KFold折叠交叉验证

文章目录 MLsklearnML流程Part2数据划分KFold折叠交叉验证Model evaluation数据划分sklearn中的模型评估demo K-fold cross-validation🎈K-foldegegeg:KFoldK-fold cross-validation ShuffleSpliteg 小结 Stratified Shuffle Splitegeg demo MLsklearnML流程Part2数…

代码随想录算法训练营第四十四天|完全背包理论基础 、518. 零钱兑换 II 、377. 组合总和 Ⅳ

文章目录 完全背包理论基础518. 零钱兑换 II377. 组合总和 Ⅳ 完全背包理论基础 纯背包问题的特点:每个物品可以无限次拿 与0-1背包唯一不同: 完全背包的物品是可以添加多次的,所以要从小到大去遍历 0-1背包不可以添加多次,需要从…

【C++】类和对象(上篇)

类和对象 面向过程和面向对象初步认识类的引入类的定义命名规范类的访问限定符及封装访问限定符封装 类的作用域类的实例化类的对象大小的计算类成员函数的this指针this指针的引出this指针的特性 面向过程和面向对象初步认识 C语言是面向过程的,关注的是过程&#…

RocketMQ-Producer

消息生产者的代码都在client模块中,相对于RocketMQ来讲,消息生产者就是客户端,也是消息的提供者。 启动流程 代码:DefaultMQProducerImpl#start public void start(final boolean startFactory) throws MQClientException {switc…

轻松掌握消息队列RabbitMQ在SpringAMQP中的实践知识点

1、介绍 spring集成了rabbitmq,可以对rabbitmq先进行安装简单了解。参考我的rabbitmq文章。 2、使用 1、基本消息队列BasicQueue案例 一个消息消费者,上个消息未处理完,队列中的消息将阻塞,导致内存泄漏 1、引入AMQP依赖 2、添…

读书:《科技论文写作与发表教程(第6版)》

科技写作是指以符合标准格式的科技论文形式在科技期刊上陈述原创性的研究。 另外,还有广义上的科技写作。 科技写作的最主要特点是表达清晰。科技写作不需要漂亮的文学修饰,要把信息清楚地传递给读者。 IMRAD格式:Introduction Methods Re…

Redis常见问题整理

一、Redis使用场景相关问题 0. 单机版Redis部署 系统环境:CentOS7 1、下载Redis所需要的镜像 yum install -y gcc tcl2、下载redis安装包 mkdir /soft cd /soft wget https://download.redis.io/releases/redis-6.2.4.tar.gz3、解压缩: tar -xvf redi…

QT-DAY3

实现ui 字体、颜色、保存文件、打开文件 #include "widget.h" #include "ui_widget.h"Widget::Widget(QWidget *parent): QWidget(parent), ui(new Ui::Widget) {ui->setupUi(this); }Widget::~Widget() {delete ui; }//字体按钮对应的槽函数 void Wid…

性能测试技术笔记:如何设计一个压测平台 ?

目录 为什么需要压测平台? 压测平台功能设计思路 压测平台技术实现方案 总结 为什么需要压测平台? 从实际工作场景出发,如果只有一两个人做性能测试工作,那其实没必要开发专门的压测平台,原因如下: 成…