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

在 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