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)方法
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 记一次.NET内存居高不下排查解决与启示
· 探究高空视频全景AR技术的实现原理
· 理解Rust引用及其生命周期标识(上)
· 浏览器原生「磁吸」效果!Anchor Positioning 锚点定位神器解析
· 没有源码,如何修改代码逻辑?
· 全程不用写代码,我用AI程序员写了一个飞机大战
· MongoDB 8.0这个新功能碉堡了,比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· DeepSeek 开源周回顾「GitHub 热点速览」
· 白话解读 Dapr 1.15:你的「微服务管家」又秀新绝活了