把数据和模型迁移到GPU: .cuda()方法和.to(device)方法
把数据从CPU迁移到GPU时,可以用.cuda()方法,也可以用.to(device)方法。示例如下。
.cuda()方法
1 2 3 4 5 6 7 | import torch import time t = time.time() b = torch.rand([ 1024 , 1024 , 10 ]) b = b.cuda() print ( 'time:{:6.3f}' . format (time.time() - t)) # 输出: time: 0.084 |
.to(device)方法 (推荐)
1 2 3 4 5 6 7 | import torch import time s = time.time() device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu' ) a = torch.rand([ 1024 , 1024 , 10 ]).to(device) print ( 'time:{:6.3f}' . format (time.time() - s)) # 输出: time: 0.087 |
.cuda()方法和.to(device)方法耗时基本差不多。
上面的方法默认把数据迁移到第0块显卡,如果要指定显卡,可以参考下面的程序:
1 2 3 4 5 6 7 8 9 10 11 | # =========== 方法一 ============== gpu = 5 torch.cuda.set_device(gpu) a = torch.rand([ 1024 , 1024 , 10 ]) a = a.cuda() # =========== 方法二 ============== gpu = 6 device = 'cuda:' + str (gpu) b = torch.randn([ 1000 , 1000 , 20 ]) b = b.to(device) |
另外,值得注意到的是,张量在GPU和CPU之间的迁移不是in-place操作,而模型在GPU和CPU之间的迁移是in-place操作。示例如下。(我这里有个疑问,in-place操作意味着存储地址没有改变,可是模型在GPU和CPU之间的迁移,存储地址怎么可能不改变?一个是GPU的显存,一个是CPU的内存,按道理迁移时存储地址是会变化的呀)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | import torch import torch.nn as nn device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu' ) x_cpu = torch.randn( 2 , 3 ) print ( 'x_cpu id: ' , id (x_cpu)) x_gpu = x_cpu.to(device) print ( 'x_gpu id: ' , id (x_gpu)) net_cpu = nn.Linear( 32 , 10 ) print ( 'net_cpu id: ' , id (net_cpu)) net_gpu = net_cpu.to(device) print ( 'net_gpu id: ' , id (net_gpu)) |
运行结果:
参考:
分类:
Python
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 记一次.NET内存居高不下排查解决与启示
· 探究高空视频全景AR技术的实现原理
· 理解Rust引用及其生命周期标识(上)
· 浏览器原生「磁吸」效果!Anchor Positioning 锚点定位神器解析
· 没有源码,如何修改代码逻辑?
· 全程不用写代码,我用AI程序员写了一个飞机大战
· MongoDB 8.0这个新功能碉堡了,比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· DeepSeek 开源周回顾「GitHub 热点速览」
· 白话解读 Dapr 1.15:你的「微服务管家」又秀新绝活了