[Pytorch学习]在GPU上训练

àì夳堔傛蜴生んèń 2021-09-28 09:46 514阅读 0赞

Pytorch学习之在GPU上训练

主要分为以下:

  1. 定义GPU

    device=torch.device(“cuda:0”if torch.cuda.is_available() else “cpu”)

  2. 将模型放到GPU中

    def Model(…):

    1. ....
    2. ....

    model=Model().to(device)

  3. 把输入和标签放到GPU中

    inputs, labels= inputs.to(device), labels.to(device)

  4. 接下来就可以正常训练拉

发表评论

表情:
评论列表 (有 0 条评论,514人围观)

还没有评论,来说两句吧...

相关阅读

    相关 Pytorch中多GPU训练指北

    前言 在数据越来越多的时代,随着模型规模参数的增多,以及数据量的不断提升,使用多GPU去训练是不可避免的事情。Pytorch在0.4.0及以后的版本中已经提供了多GPU训