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() 
相关推荐
love530love2 分钟前
Docker 稳定运行与存储优化全攻略(含可视化指南)
运维·人工智能·windows·docker·容器
剑桥折刀s11 分钟前
Python打卡:Day46
python
HeartException28 分钟前
量子计算+AI芯片:光子计算如何重构神经网络硬件生态
人工智能
摸鱼仙人~32 分钟前
Minstrel:多智能体协作生成结构化 LangGPT 提示词
人工智能·提示词
巴里巴气33 分钟前
Python爬虫图片验证码和滑块验证码识别总结
爬虫·python
sword devil9001 小时前
PYQT实战:智能家居中控
python·智能家居·pyqt
NetX行者1 小时前
FastMCP:用于构建MCP服务器的开源Python框架
服务器·python·开源
超龄超能程序猿1 小时前
(3)机器学习小白入门 YOLOv: 解锁图片分类新技能
python·numpy·pandas·scipy
AI街潜水的八角2 小时前
深度学习图像分类数据集—濒危动物识别分类
人工智能·深度学习
waynaqua2 小时前
FastAPI开发AI应用一:实现连续多轮对话
python·openai