CPU和GPU张量类型转换

Small_Lemon_Tree

于 2020-07-10 16:57:36 发布

阅读量1.3k

点赞数 1

分类专栏： Pytorch

版权

9 篇文章

订阅专栏

本文介绍在PyTorch中如何实现CPU张量到GPU张量的转换，以及相反操作，通过调用Tensor.cuda()和Tensor.cpu()函数。文章还提供了一个实用的函数commonType(Tensor)，用于自动判断并执行张量的设备转换，确保代码的兼容性和效率。

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

(1) CPU张量 ----> GPU张量，使用Tensor.cuda()

(2) GPU张量 ----> CPU张量，使用Tensor.cpu()

我们可以通过torch.cuda.is_available()函数来判断当前的环境是否支持GPU，如果支持，则返回GPU类型，否则返回CPU类型。

def commonType(Tensor):
    if torch.cuda.is_available():
        cuda = "cuda:0"
        return Tensor.cuda(cuda)
    else:
        return Tensor

更多：https://blog.youkuaiyun.com/HowardWood/article/details/79508648