torch使用GPU加速

最新推荐文章于 2024-09-08 12:16:53 发布

他是一个俗人啊

最新推荐文章于 2024-09-08 12:16:53 发布

阅读量1.1k

点赞数

分类专栏：深度学习文章标签：深度学习 pytorch python

本文链接：https://blog.csdn.net/weixin_44078014/article/details/131980381

版权

4 篇文章 0 订阅

订阅专栏

文章介绍了如何下载CUDA和torch库，以及如何确保torch与Python版本匹配以使用GPU加速。内容包括将模型、函数和数据迁移到GPU上，如自定义模型需继承torch.nn.Module，模型和数据的cuda()函数调用，以及GPU上loss函数的计算。

摘要由CSDN通过智能技术生成

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

将模型，函数，数据都放入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()

关注

专栏目录