首页 文章资讯内容详情

如何将 Torch Tensor 从 CPU 移动到 GPU,反之亦然?

2026-06-02 1 花语

在CPU上定义的火炬张量可以移动到GPU,反之亦然。对于高维张量计算,GPU利用并行计算的能力来减少计算时间。

像图像这样的高维张量是高度计算密集型的,如果在CPU上运行会花费太多时间。因此,我们需要将此类张量移至GPU。

语法

要将Torch张量从CPU移动到GPU,使用以下语法/es-

Tensor.to("cuda:0") or Tensor.to(cuda)

和,

Tensor.cuda()

要将火炬张量从GPU移动到CPU,使用以下语法-

Tensor.to("cpu")

和,

Tensor.cpu()

让我们举几个例子来演示张量如何从CPU移动到GPU,反之亦然。

注意-我为每个程序提供了两个不同的输出。一个输出用于仅具有CPU的系统,另一个输出用于具有GPU和CPU的系统。

示例1

# Python program to move a tensor from CPU to GPU # import torch library import torch # create a tensor x = torch.tensor([1.0,2.0,3.0,4.0]) print("Tensor:", x) # check tensor device (cpu/cuda) print("张量装置:", x.device) # Move tensor from CPU to GPU # check CUDA GPU is available or not print("CUDA图形处理器:", torch.cuda.is_available()) if torch.cuda.is_available(): x = x.to("cuda:0") # or x=x.to("cuda") print(x) # now check the tensor device print("张量装置:", x.device)

输出1-当GPU不可用时

Tensor: tensor([1., 2., 3., 4.]) 张量装置: cpu CUDA图形处理器: False tensor([1., 2., 3., 4.]) 张量装置: cpu

输出2-当GPU可用时

Tensor: tensor([1., 2., 3., 4.]) 张量装置: cpu CUDA图形处理器: True tensor([1., 2., 3., 4.], device=cuda:0) 张量装置: cuda:0

示例2

# Python program to move a tensor from CPU to GPU # import torch library import torch # create a tensor on CPU x = torch.tensor([1.0,2.0,3.0,4.0]) print("Tensor:", x) print("张量装置:", x.device) # Move tensor from CPU to GPU if torch.cuda.is_available(): x = x.cuda() print(x) # now check the tensor device print("张量装置:", x.device)

输出1-如果GPU不可用

Tensor: tensor([1., 2., 3., 4.]) 张量装置: cpu tensor([1., 2., 3., 4.]) 张量装置: cpu

输出2-如果GPU可用

Tensor: tensor([1., 2., 3., 4.]) 张量装置: cpu tensor([1., 2., 3., 4.], device=cuda:0) 张量装置: cuda:0

示例3

# Python program to move a tensor from GPU to CPU # import torch library import torch # create a tensor on GPU if torch.cuda.is_available(): x = torch.tensor([1.0,2.0,3.0,4.0], device = "cuda") print("Tensor:", x) print("张量装置:", x.device) # Move tensor from GPU to CPU x = x.to("cpu") # x = x.cpu() print(x) # Now check the tensor device print("张量装置:", x.device)

输出1-如果GPU不可用

Tensor: tensor([1., 2., 3., 4.]) 张量装置: cpu tensor([1., 2., 3., 4.]) 张量装置: cpu

输出2-如果GPU可用

Tensor: tensor([1., 2., 3., 4.], device=cuda:0) 张量装置: cuda:0 tensor([1., 2., 3., 4.]) 张量装置: cpu