当前位置：首页 > article >正文

pytorch 张量创建基础

article 2025/2/1 17:41:34

在默认情况下，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")