pytorch 张量创建基础

发布于:2025-02-10 ⋅ 阅读:(30) ⋅ 点赞:(0)

在默认情况下,torch.tensor 会在 CPU 上创建张量。以下是一个简单的示例:

import torch

# 使用 torch.tensor 创建一个张量
tensor_cpu = torch.tensor([1, 2, 3])
print("张量所在设备:", tensor_cpu.device)

在上述代码中,通过 torch.tensor 创建了一个张量 tensor_cpu,并使用 tensor_cpu.device 查看该张量所在的设备,输出结果通常为 cpu,这表明张量是在 CPU 上创建的。

将张量移动到 GPU
如果你希望将 torch.tensor 创建的张量放到 GPU 上,可以使用以下几种方法:

1 在创建时指定设备

import torch

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    # 在 GPU 上创建张量
    tensor_gpu = torch.tensor([1, 2, 3], device='cuda')
    print("张量所在设备:", tensor_gpu.device)
else:
    print("没有可用的 GPU")

2 创建后移动到 GPU

# 在 CPU 上创建张量
tensor_cpu = torch.tensor([1, 2, 3])

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    # 使用 to 方法将张量移动到 GPU
    tensor_gpu_1 = tensor_cpu.to('cuda')
    print("使用 to 方法移动后张量所在设备:", tensor_gpu_1.device)

    # 使用 cuda 方法将张量移动到 GPU
    tensor_gpu_2 = tensor_cpu.cuda()
    print("使用 cuda 方法移动后张量所在设备:", tensor_gpu_2.device)
else:
    print("没有可用的 GPU")

多 GPU 情况

import torch

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    # 在第二个 GPU 上创建张量(GPU 编号从 0 开始)
    tensor_gpu = torch.tensor([1, 2, 3], device='cuda:1')
    print("张量所在设备:", tensor_gpu.device)
else:
    print("没有可用的 GPU")

在这段代码中,使用 device=‘cuda:1’ 将张量创建在了编号为 1 的 GPU 上。
综上所述,torch.tensor 默认在 CPU 上创建张量,但可以通过指定 device 参数或使用 to、cuda 方法将张量放到 GPU 上。