PyTorch|在张量运算中使用GPU

霜溪2024-01-07 20:37

pytorch在张量运算时允许我们在GPU上进行计算，我们可以采用这些方法将数据无缝的移入或移出GPU。

当我们进入GPU时，我们可以使用cuda()方法 ，当我们进入CPU时，我们可以使用**cpu()**方法。

同时，我们还可以使用to() 方法，去GPU的时候，我们写到**('cuda')** ，去CPU的时候，我们写到**('cpu')**。

在训练过程中，要想利用我们的GPU，有两个基本要求：

数据必须移到GPU上
网络必须移到GPU上

默认情况下，在创建PyTorch张量或PyTorch神经网络模块时，会在CPU上初始化相应的数据。具体来说，这些数据存在于CPU的内存中。

下面是一些实例：

复制代码

>>> import torch
>>> data=torch.tensor([1,2,3,4])>>> data.devicedevice(type='cpu')
>>> v1=data.to('cuda')>>> v1.devicedevice(type='cuda', index=0)>>> v2=data.cuda()>>> v2.devicedevice(type='cuda', index=0)
>>> v3=v2.to('cpu')>>> v3.devicedevice(type='cpu')>>> v4=v2.cpu()>>> v4.device

然而，代码需要在不同的设备都能运行，也就是在不修改代码的情况下在GPU上可以使用或者在只能使用CPU的设备上运行。

例如，假设我们写的代码到处使用cuda()方法，然后，我们把代码交给一个没有GPU的用户，这样做是行不通的。

对此，我们可以这样处理：

复制代码

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")input=data.to(device)network=Network()net=network.to(device)

总结起来，在我们编写的代码中，尽量使用这种形式：

复制代码

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")***=***.to(device)