torch使用GPU加速

下载cuda

下载torch

下载与python版本适配的GPU版torch

使用GPU加速

将模型,函数,数据都放入GPU中

  • 将模型加入GPU
    # 自定义模型要继承torch.nn.Module才能使用GPU加速
    # 在模型后加cuda函数
    model = TransE(num_nodes, num_relations, hidden_channels).cuda()
    
  • 将函数放入GPU,例外:模型model的__init__函数中出现的self调用的函数,数据将自动放入GPU(模型已被放入GPU)
    # 例如将loss函数放入GPU
    l = train_model.loss(head_index.cuda(), relation_type.cuda(), tail_index.cuda()).cuda()
    
  • 将数据放入GPU
    # 基本将tensor类型的数据都放入GPU
    head_index.cuda()
    

你可能感兴趣的:(深度学习,深度学习,pytorch,python)