PyTorch查看模型和数据是否在GPU上

模型和数据可以在CPU和GPU上来回迁移,怎么判断模型和数据在哪里呢?

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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

此外,用.is_cuda也可以判断模型和数据是否在GPU上,例如: data.is_cuda 

模型和数据在GPU和CPU之间进行迁移的方法,更具体的可以参考:把数据和模型迁移到GPU: .cuda()方法和.to(device)方法 

 

posted @   Picassooo  阅读(33752)  评论(0编辑  收藏  举报
编辑推荐:
· 记一次.NET内存居高不下排查解决与启示
· 探究高空视频全景AR技术的实现原理
· 理解Rust引用及其生命周期标识(上)
· 浏览器原生「磁吸」效果!Anchor Positioning 锚点定位神器解析
· 没有源码,如何修改代码逻辑?
阅读排行:
· 全程不用写代码,我用AI程序员写了一个飞机大战
· MongoDB 8.0这个新功能碉堡了,比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· DeepSeek 开源周回顾「GitHub 热点速览」
· 白话解读 Dapr 1.15:你的「微服务管家」又秀新绝活了
点击右上角即可分享
微信分享提示