PyTorch_指定运算设备 (包含安装 GPU 的 PyTorch)

PyTorch默认会将张量创建在 CPU 控制的内存中,即:默认的运算设备为 CPU。我们也可以将张量创建在 GPU 上,能够利用对于矩阵计算的优势加快模型训练。

将张量移动到 GPU 上有两种方法:

  1. 使用 cuda 方法
  2. 直接在 GPU 上创建张量
  3. 使用 to 方法指定设备

安装含有 GPU 的 PyTorch

通过这个可以判断电脑里是否已经安装 CUDA

python 复制代码
import torch 

print(torch.cuda.is_available()) # 判断是否有可用的 GPU 设备

如果结果输出是 False,说明设备里没有 GPU 的 PyTorch 的版本。

可以通过 PyTorch官网来安装含有 GPU 的 PyTorch 的版本。

安装 CUDA 12.8 版本的pip命令

复制代码
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

这样就安装好含有 GPU 的 PyTorch 了。


代码

python 复制代码
import torch 

# 使用 cuda 方法
def test01():
    data = torch.tensor([10, 20, 30])
    print("存储设备:", data.device)

    # 将张量移动到 GPU 设备上
    data = data.cuda()
    print("存储设备:", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.cpu()
    print("存储设备:", data.device)

# 直接将张量创建在指定设备上
def test02():
    data = torch.tensor([10, 20, 30], device='cuda')
    print("存储设备:", data.device)

    # 将张量移动到 CPU 设备上
    data = data.cpu()
    print("存储设备:", data.device)

# 使用 to 方法
def test03():
    data = torch.tensor([10, 20, 30])
    print("存储设备:", data.device)

    # 将张量移动到 GPU 设备上
    data = data.to('cuda')
    print("存储设备:", data.device)
    
    # 将张量从 GPU 移动到 CPU 设备上
    data = data.to('cpu')
    print("存储设备:", data.device)


# 注意点:张量存储在不同设备上的张量不能够直接运算
def test04():
    data1 = torch.tensor([10,20,30])
    data2 = torch.tensor([1,2,3], device='cuda')

    # RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
    # data1 = data1.to('cuda') # 这样可以解决了
    data = data1 + data2 
    print(data)

    # 如果你的电脑上安装 pytorch 不是 gpu 版本的,或者电脑本身没有 gpu 设备环境
    # 否则下面的调用 cuda 函数的代码会报错
    # data1 = data1.cuda() 

if __name__ == "__main__":
    test04() 
相关推荐
点云SLAM13 分钟前
PyTorch 中.backward() 详解使用
人工智能·pytorch·python·深度学习·算法·机器学习·机器人
androidstarjack19 分钟前
波士顿动力给机器人装上AI大脑,人类故意使绊子也不怕了!
人工智能·机器人
B1118521Y4642 分钟前
flask的使用
后端·python·flask
Learn Beyond Limits1 小时前
Transfer Learning|迁移学习
人工智能·python·深度学习·神经网络·机器学习·ai·吴恩达
程序员三明治1 小时前
三、神经网络
人工智能·深度学习·神经网络
hundaxxx3 小时前
自演化大语言模型的技术背景
人工智能
数智顾问3 小时前
【73页PPT】美的简单高效的管理逻辑(附下载方式)
大数据·人工智能·产品运营
love530love3 小时前
【保姆级教程】阿里 Wan2.1-T2V-14B 模型本地部署全流程:从环境配置到视频生成(附避坑指南)
人工智能·windows·python·开源·大模型·github·音视频
木头左3 小时前
结合机器学习的Backtrader跨市场交易策略研究
人工智能·机器学习·kotlin
Coovally AI模型快速验证3 小时前
3D目标跟踪重磅突破!TrackAny3D实现「类别无关」统一建模,多项SOTA达成!
人工智能·yolo·机器学习·3d·目标跟踪·无人机·cocos2d