使用PyTorch将数据从CPU移动到GPU的四个方法
作者:阿罗的小小仓库
问题:已知一个张量cpu_tensor,将其移动到GPU
import torch # 在 CPU 上创建一个张量 cpu_tensor = torch.tensor([1, 2, 3])
一、使用 .to() 方法:
# 将张量移动到 GPU gpu_tensor = cpu_tensor.to('cuda') print("CPU Tensor:", cpu_tensor) print("GPU Tensor:", gpu_tensor)
.to() 方法,直接在原始张量上进行设备转移。它不会创建新的张量,而是在原地更新了原始张量的设备信息。这样的方式很高效,避免了不必要的数据拷贝。
二、使用 .cuda() 方法:
# 将张量移动到 GPU gpu_tensor = cpu_tensor.cuda() print("CPU Tensor:", cpu_tensor) print("GPU Tensor:", gpu_tensor)
三、使用 torch.Tensor 构造函数:
# 使用构造函数将张量移动到 GPU gpu_tensor = torch.Tensor(cpu_tensor).cuda() print("CPU Tensor:", cpu_tensor) print("GPU Tensor:", gpu_tensor)
四、使用 torch.tensor 构造函数:
# 使用构造函数将张量移动到 GPU gpu_tensor = torch.tensor(cpu_tensor, device='cuda') print("CPU Tensor:", cpu_tensor) print("GPU Tensor:", gpu_tensor)
四个方法的输出一致:
CPU Tensor: tensor([1,2,3])
GPU Tensor: tensor([1,2,3],device='cuda:0')
总结:
.to() 方法是一个通用的设备转移方法,不仅可以用于 CPU 和 GPU 之间的转移,还可以在不同的 GPU 之间转移。
除了设备,.to() 方法还可以指定数据类型(dtype),如 .to('cuda', dtype=torch.float32)。
返回的是新的张量,原始张量不会被修改。
.to() 方法的语法为:tensor.to(device, dtype=None, non_blocking=False)。
.cuda() 方法是 .to('cuda') 的一种快捷方式,专门用于将张量移动到 GPU。
不支持在不同的 GPU 之间转移。
与 .to() 方法不同,.cuda() 不接受 dtype 参数。
返回的是新的张量,原始张量不会被修改。
torch.Tensor 构造函数,先用torch.Tensor(cpu_tensor) 创建了一个新的张量,该张量的数据与 cpu_tensor 共享。然后,.cuda() 方法在原地将该张量的设备信息更新为 GPU。这个方法可以用于在已经存在的张量上执行设备转移。
torch.tensor 构造函数,并通过 device 参数指定了目标设备为 GPU。它实际上是在创建一个新的张量,将原始张量的数据复制到了 GPU 上。这个方法用于在创建张量时就指定设备。
这些方法的选择取决于个人偏好和代码上下文。在实际使用中,通常使用 .to('cuda') ,这种方式最为简洁和通用。
以上就是使用PyTorch将数据从CPU移动到GPU的四个方法的详细内容,更多关于PyTorch数据从CPU移动到GPU的资料请关注脚本之家其它相关文章!