本文介绍如何使用PyTorch将数据加载到GPU中进行高效运算。通过简单的代码实现，可以在训练前将Tensor变量复制到指定的GPU设备上，确保后续计算都在GPU环境中完成。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
这段代码一般写在读取数据之前，torch.device代表将torch.Tensor分配到的设备的对象。torch.device包含一个设备类型（‘cpu’或‘cuda’）和可选的设备序号。如果设备序号不存在，则为当前设备。如：torch.Tensor用设备构建‘cuda’的结果等同于‘cuda：X’，其中X是torch.cuda.current_device()的结果。
mytensor = my_tensor.to(device)
这行代码的意思是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去，之后的运算都在GPU上进行。
这句话需要写的次数等于需要保存GPU上的tensor变量的个数；一般情况下这些tensor变量都是最开始读数据时的tensor变量，后面衍生的变量自然也都在GPU上
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
