【PyTorch】PyTorch深度学习框架实战(二):torchrun

2024-08-13 16:28:30 浏览数 (3)

一、引言

  • PyTorch由facebook人工智能研究院研发,2017年1月被提出,是一个开源的Python机器学习库,基于Torch,用于自然语言处理等应用程序。PyTorch既可以看作加入了GPU支持的numpy,同时也可以看成一个拥有自动求导功能的强大的深度神经网络。
  • ​PyTorch的前身是Torch,其底层和Torch框架一样,但是使用Python重新写了很多内容,不仅更加灵活,支持动态图,而且提供了Python接口。它是由Torch7团队开发,是一个以Python优先的深度学习框架,不仅能够实现强大的GPU加速,同时还支持动态神经网络。

二、PyTorch介绍

2.1 PyTorch特点

PyTorch是一个python包,提供两个高级功能:

2.1.1 支持GPU加速的张量计算库

张量(tensor):可以理解为多位数组,是Pytorch的基本计算单元,Pytorch的特性就是可以基于GPU快速完成张量的计算,包括求导、切片、索引、数学运算、线性代数、归约等

代码语言:javascript复制
import torch
import torch.nn.functional as F

# 1. 张量的创建
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
y = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(x) #tensor([[1, 2, 3],[4, 5, 6]])
print(y) #tensor([[1, 2, 3],[4, 5, 6]])

# 2. 张量的运算
z=x y
print(z) #tensor([[2, 4, 6],[8, 10, 12]])

# 3. 张量的自动求导
x = torch.tensor(3.0, requires_grad=True)
print(x.grad) #None

y = x**2 
y.backward()
print(x.grad) #tensor(6.)
2.1.2 包含自动求导系统的动态图机制

Pytorch提供了一种独一无二的构建神经网络的方式:动态图机制

不同于TensorFlow、Caffe、CNTK等静态神经网络:网络构建一次反复使用,如果修改了网络不得不重头开始。

在Pytorch中,使用了一种“反向模式自动微分的技术(reverse-mode auto-differentiation)”,允许在零延时或开销的情况下任意更改网络。

2.2 PyTorch安装

这里建议大家采用conda创建环境,采用pip管理pytorch包

1.建立名为pytrain,python版本为3.11的conda环境

代码语言:javascript复制
conda create -n pytrain python=3.11
conda activate pytrain

2.采用pip下载torch和torchvision包

代码语言:javascript复制
pip install torch  torchvision torchmetrics  -i https://mirrors.cloud.tencent.com/pypi/simple

这里未指定版本,默认下载最新版本torch-2.3.0、torchvision-0.18.0以及其他一堆依赖。

三、PyTorch之分布式训练框架torchrun

3.1 torchrun简介

torchrun是PyTorch库中用于启动分布式训练的命令行工具,特别是在使用PyTorch Distributed Package时。它简化了分布式训练的启动过程,自动处理了如初始化进程群、设置环境变量等复杂步骤,使得在多GPU或者多节点环境下的分布式训练变得更加便捷

3.2 torchrun主要用途

  • 多GPU训练:在单机多GPU环境下执行分布式训练。
  • 多节点训练:在多台机器的多GPU环境下执行分布式训练,通过网络连接。
  • 自动环境配置:自动设置MASTER_ADDR和MASTER_PORT等环境变量,用于进程间的通信。
  • 进程管理:根据指定的后端(如NCCL、Gloo)和运行策略(如单进程多线程、多进程)来启动和管理训练进程。

3.3 torchrun基本用法

代码语言:javascript复制
torchrun [OPTIONS] SCRIPT [SCRIPT_ARGS]

-n, --nproc_per_node: 指定每个节点上的GPU数量或进程数。
--nnodes: 指定参与训练的节点总数。
--node_rank: 当在多个节点上运行时,指定当前节点的排名。
--master_addr: 指定主节点的IP地址。
--master_port: 指定主节点的端口号。
--no_spawn: 在单节点上不使用多进程,直接运行,主要用于调试。

假设你有一个训练脚本train.py,想要在本地机器的4个GPU上进行分布式训练,可以使用以下命令

代码语言:javascript复制
torchrun --nproc_per_node=4 train.py

如果是在多节点环境下,比如有两个节点,每个节点有4个GPU,你可能需要在每个节点上分别执行(假设主节点IP为192.168.1.100):

在节点1(假设其IP也是192.168.1.100,作为主节点)上:

代码语言:javascript复制
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.100 --master_port=29500 train.py

在节点2上:

代码语言:javascript复制
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr=192.168.1.100 --master_port=29500 train.py

使用时需要确保网络配置允许节点间的通信,并且所有节点上都安装了PyTorch和必要的依赖。

四、总结

本文首先对PyTorch的历史、特点和安装方法进行简要介绍,接着对PyTorch的分布式训练框架torchrun实用方法进行整理,欢迎大家点赞收藏,以备查阅。

0 人点赞