PyTorch查看模型和数据是否在GPU上
2023-02-18 16:32:57 时间
Example
import torch
import torch.nn as nn
# ----------- 判断模型是在CPU还是GPU上 ----------------------
model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)
print(next(model.parameters()).device) # 输出:cpu
model = model.cuda()
print(next(model.parameters()).device) # 输出:cuda:0
model = model.cpu()
print(next(model.parameters()).device) # 输出:cpu
# ----------- 判断数据是在CPU还是GPU上 ----------------------
data = torch.ones([2, 3])
print(data.device) # 输出:cpu
data = data.cuda()
print(data.device) # 输出:cuda:0
data = data.cpu()
print(data.device) # 输出:cpu
相关文章
- Netty 概述
- Oracle 伪列(ROWNUM)机制
- mysql只会使用到一个索引的原因
- SQL优化-三
- SQL优化-二
- SQL优化-一
- RedisUtil-redisTemplate-setNX
- 数据库无限层级分类设计
- 魔方
- CountDownLatch在SpringBoot中配合@Async使用
- 会话刷新Token校验流程
- Mybatis 夺命十八问,顶不住了! - 里奥ii的文章 - 知乎
- 阿里架构演变
- 自定义注解和事务优先级
- MyCat 连接异常 解决参考链接
- SpringBoot手动事务参考链接
- SimpleDateFormat并发隐患及其解决
- 框架收藏
- 事务隔离级别
- 循环中使用spring事务