PyTorch是一个python包,提供两个高级功能:
张量(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.)
Pytorch提供了一种独一无二的构建神经网络的方式:动态图机制
不同于TensorFlow、Caffe、CNTK等静态神经网络:网络构建一次反复使用,如果修改了网络不得不重头开始。
在Pytorch中,使用了一种“反向模式自动微分的技术(reverse-mode auto-differentiation)”,允许在零延时或开销的情况下任意更改网络。
这里建议大家采用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以及其他一堆依赖。
torchrun是PyTorch库中用于启动分布式训练的命令行工具,特别是在使用PyTorch Distributed Package时。它简化了分布式训练的启动过程,自动处理了如初始化进程群、设置环境变量等复杂步骤,使得在多GPU或者多节点环境下的分布式训练变得更加便捷
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实用方法进行整理,欢迎大家点赞收藏,以备查阅。