PyTorch使用GPU进行Tensor及模型计算

文章目录

  • [1. 计算设备:GPU/CPU](#1. 计算设备:GPU/CPU)
  • [2. Tensor的GPU计算](#2. Tensor的GPU计算)
  • [3. 模型的GPU计算](#3. 模型的GPU计算)

对复杂的神经网络和大规模的数据来说,使用CPU来计算可能不够高效。这里,我们将介绍如何使用单块NVIDIA GPU来计算。

首先,需要确保已经安装好了PyTorch GPU版本,具体操作步骤可以参考博客:Ubuntu 22.04安装cuda及Pytorch教程

查看显卡占用情况:

bash 复制代码
nvidia-smi

每隔2s刷新一次显卡占用情况:

python 复制代码
watch -n 2 nvidia-smi

输出如下:

可以看到我这里只有两块 NVIDIA GeForce RTX 2080,单块显卡有11G的显存。

1. 计算设备:GPU/CPU

PyTorch可以指定用来存储和计算的设备,如使用内存的CPU或者使用显存的GPU。默认情况下,PyTorch会将数据创建在内存,然后利用CPU来计算。

  1. 查看GPU是否可用:
python 复制代码
torch.cuda.is_available() # 输出 True
  1. 查看GPU数量:
python 复制代码
torch.cuda.device_count() # 输出 2
  1. 查看当前GPU索引号,索引号从0开始:
python 复制代码
torch.cuda.current_device() # 输出 0
  1. 根据索引号查看GPU名字:
python 复制代码
torch.cuda.get_device_name(0) # 输出 'GeForce RTX 2080'

2. Tensor的GPU计算

默认情况下,Tensor会被存在内存上(即CPU中)。因此,之前我们每次打印Tensor的时候看不到GPU相关标识。

python 复制代码
x = torch.tensor([1, 2, 3])
x

输出:

python 复制代码
tensor([1, 2, 3])

1. 使用.cuda()可以将CPU上的Tensor转换(复制)到GPU上:

如果有多块GPU,我们用.cuda(i)来表示第 i i i 块GPU及相应的显存( i i i从0开始)且cuda(0)和cuda()等价。

python 复制代码
x = x.cuda(0)
x

输出:

python 复制代码
tensor([1, 2, 3], device='cuda:0')

2. 通过Tensor的device属性来查看该Tensor所在的设备:

python 复制代码
x.device

输出:

python 复制代码
device(type='cuda', index=0)

3. 创建Tensor时指定设备:

我们也可以直接在创建Tensor时就指定设备。

python 复制代码
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = torch.tensor([1, 2, 3], device=device)
# or
x = torch.tensor([1, 2, 3]).to(device)
x

输出:

python 复制代码
tensor([1, 2, 3], device='cuda:0')

4. 如果对在GPU上的数据进行运算,那么结果还是存放在GPU上。

python 复制代码
y = x**2
y

输出:

python 复制代码
tensor([1, 4, 9], device='cuda:0')

5. GPU与CPU无法混用。

需要注意的是,存储在不同位置中的数据是不可以直接进行计算的。即存放在CPU上的数据不可以直接与存放在GPU上的数据进行运算,位于不同GPU上的数据也是不能直接进行计算的。

python 复制代码
z = y + x.cpu()

会报错:

python 复制代码
RuntimeError: Expected object of type torch.cuda.LongTensor but found type torch.LongTensor for argument #3 'other'

3. 模型的GPU计算

同Tensor类似,PyTorch模型也可以通过.cuda转换到GPU上。我们可以通过检查模型的参数的device属性来查看存放模型的设备。

python 复制代码
net = nn.Linear(3, 1)
list(net.parameters())[0].device

输出:

python 复制代码
device(type='cpu')

可见模型在CPU上。

1. 将模型转换到GPU上:

python 复制代码
net.cuda()
list(net.parameters())[0].device

输出:

python 复制代码
device(type='cuda', index=0)

2. 模型输入的Tensor和模型需在同一设备上。

同样的,我么需要保证模型输入的Tensor和模型都在同一设备上,否则会报错。

这里,我们给出两种将Tensor和模型转换到GPU设备上的写法:

python 复制代码
import torch

x = torch.rand(2, 3)
net = nn.Linear(3, 1)
  • 写法1:基于 cuda() 函数
python 复制代码
x = x.cuda()
net = net.cuda()
net(x)
  • 写法2:基于 to(device) 函数 (推荐)
python 复制代码
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = x.to(device)
net = net.to(device)
net(x)

to(device) 这种写法可以自动化获取设备(CPU或GPU),并将Tensor和Model自动转化到设备上,在实际使用过程中更为常见。

  1. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 自动获取设备。
  2. x = x.to(device) 基于 to()函数将张量数据转换到GPU设备上。
  3. net = net.to(device) 基于 to()函数将模型转换到GPU设备上。

输出:

python 复制代码
tensor([[-0.5800],
        [-0.2995]], device='cuda:0', grad_fn=<ThAddmmBackward>)
相关推荐
飞哥数智坊9 小时前
从CodeBuddy翻车到MasterGo救场,我的小程序UI终于焕然一新
人工智能
AKAMAI11 小时前
跳过复杂环节:Akamai应用平台让Kubernetes生产就绪——现已正式发布
人工智能·云原生·云计算
databook13 小时前
Manim实现闪光轨迹特效
后端·python·动效
新智元13 小时前
阿里王牌 Agent 横扫 SOTA,全栈开源力压 OpenAI!博士级难题一键搞定
人工智能·openai
新智元13 小时前
刚刚,OpenAI/Gemini 共斩 ICPC 2025 金牌!OpenAI 满分碾压横扫全场
人工智能·openai
机器之心13 小时前
OneSearch,揭开快手电商搜索「一步到位」的秘技
人工智能·openai
阿里云大数据AI技术13 小时前
2025云栖大会·大数据AI参会攻略请查收!
大数据·人工智能
Juchecar14 小时前
解惑:NumPy 中 ndarray.ndim 到底是什么?
python
YourKing14 小时前
yolov11n.onnx格式模型转换与图像推理
人工智能
sans_14 小时前
NCCL的用户缓冲区注册
人工智能