【PyTorch】PyTorch深度学习框架实战（二）：torchrun

一、引言

• 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快速完成张量的计算，包括求导、切片、索引、数学运算、线性代数、归约等

​

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环境

conda create -n pytrain python=3.11
conda activate pytrain

2.采用pip下载torch和torchvision包

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基本用法

torchrun [OPTIONS] SCRIPT [SCRIPT_ARGS]

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

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

torchrun --nproc_per_node=4 train.py

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

在节点1（假设其IP也是192.168.1.100，作为主节点）上：

torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=192.168.1.100 --master_port=29500 train.py

在节点2上：

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实用方法进行整理，欢迎大家点赞收藏，以备查阅。