GPU能大幅提升计算速度,为了达到GPU加速训练的目的,就要将进行运算的数据传入GPU,回想一下神经网络,所有的数据不外乎是模型参数和数据集,所以我们只要将模型参数和数据集传入GPU即可。
用 .cuda() 可以达到这个目的

class Net(torch.nn.Module):
      def __init__(self):...   #我只是省略了网络详细内容 0.0
      def forward(self, x):...
      return x
my_cnn=Net().cuda()   #将网络parameters传入GPU

数据集同理,读取后后边加上 .cuda() 即可

注意GPU里的数据和CPU里的是不能混用, 需要将数据转成CPU的话,后边加上 .cpu() 即可。

附上几个常用的GPU指令:

torch.cuda.is_available()
cuda是否可用;

torch.cuda.device_count()
返回GPU数量;

torch.cuda.get_device_name(0)
返回GPU名字,设备索引默认从0开始;

torch.cuda.current_device()
返回当前设备索引

torch.cuda.empty_cache()
释放缓存分配器当前持有的所有未占用的缓存显存

了解更多:https://pytorch.apachecn.org/docs/1.2/cuda.html

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐