pytorch指定gpu运行

最新推荐文章于 2026-05-12 05:25:18 发布

原创最新推荐文章于 2026-05-12 05:25:18 发布 · 2.5k 阅读

3 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #gpu

收录于

pytorch学习专栏收录该内容

7 篇文章

订阅专栏

博客介绍了在使用PyTorch时，利用GPU的操作方式，包括直接在终端中操作，以及在脚本中指定和操作，为相关开发提供了使用GPU的途径。

直接在终端中

 指定编号为1(默认从0开始)的GPU：           CUDA_VISIBLE_DEVICES=1 python xxx.py

在脚本中指定

 					os.environ['CUDA_VISIBLE_DEVICES'] = gpus             #gpus可以以数字的形式或者列表的形式

在脚本中操作

cuda = torch.device('cuda')     # 默认CUDA设备（比如上一条的脚本指定gpu，或者在torch.cuda.device()里面等等情况）
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2')  # GPU 2 (these are 0-indexed)

#通过tensor.to()和tensor.cuda()的方式把tensor转移到特定设备
a = torch.randn(2).to(cuda2)     #把tensor生成在GPU2上面
b = torch.randn(2).cuda(cuda2)