pytorch-设备间转换

lakomi

于 2020-08-03 17:59:02 发布

阅读量861

点赞数

分类专栏： pytorch 文章标签： gpu cuda

本文链接：https://blog.youkuaiyun.com/Q_s_qiu/article/details/107770143

版权

pytorch 专栏收录该内容

8 篇文章

订阅专栏

本文介绍了如何将PyTorch中的Tensor从CPU移动到GPU，包括使用cuda()和to()方法的具体操作流程。同时，文章还讨论了在多GPU环境下，如何指定特定的GPU设备进行计算。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

移动Tensor到GPU

虽然现在不支持CUDA，先存起，没准后边就用上了

注意，将Tensor移动到GPU，这一步操作需要CUDA设备支持

一个GPU

cpu_a = torch.rand(4, 3)
print(cpu_a.type())  # torch.FloatTensor

# 如果不支持CUDA，此处将会报错。
# 将tensor移动到gpu
gpu_a = cpu_a.cuda()
print(gpu_a.type())  # torch.cuda.FloatTensor

# 将tensor移动到cpu
cpu_b = gpu_a.cpu()
print(cpu_b.type())  # torch.FloatTensor

多个GPU

# 若有多GPU，使用to()方法 指明使用哪个设备
# 确定是否有cuda设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
# 将tensor传送到设备
gpu_b = cpu_b.to(device)
print(gpu_b.type())