电脑基础 · 2023年3月31日

torch.cuda常用指令

CUDA(Compute Unified Device Architecture),是显卡厂商NVIDIA推出的运算平台。 CUDA™是一种由NVIDIA推出的通用并行计算架构,该架构使GPU能够解决复杂的计算问题。 它包含了CUDA指令集架构(ISA)以及GPU内部的并行计算引擎。 开发人员可以使用C语言来为CUDA™架构编写程序,所编写出的程序可以在支持CUDA™的处理器上以超高性能运行。

Pytorch通过cuda指令允许让模型、数据加载到GPU上,常用指令如下:

1. torch.cuda.is_available()

cuda是否可用  

import torch
print(torch.cuda.is_available())

torch.cuda常用指令

2. torch.cuda.device_count()

查看GPU数量

import torch
print(torch.cuda.device_count())

torch.cuda常用指令

3. torch.cuda.get_device_name()

查看DEVICE(GPU)名

import torch
print(torch.cuda.get_device_name())

torch.cuda常用指令

4. torch.cuda.current_device()

检查目前使用GPU的序号

import torch
print(torch.cuda.current_device())

torch.cuda常用指令

5. torch.cuda.set_device()

指定使用的卡

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #多卡

只指定主显卡,如下:

import torch
torch.cuda.set_device(1)
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

torch.cuda常用指令

指定特定显卡,如下:

import torch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
torch.cuda.set_device(1)
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

torch.cuda常用指令

使用torch.cuda.set_device()可以更方便地将模型和数据加载到对应GPU上, 在定义模型之前加入一行代码即可,但是这种写法的优先级低,如果model.cuda()中指定了参数,那么torch.cuda.set_device()会失效,而且pytorch的官方文档中明确说明,不建议用户使用该方法。

6.  .cuda()

指定模型和数据加载到对应的GPU,以net.cuda()为例,加载方法为:

net.cuda(gpu_id) # 输入参数为int类型,只能指定一张显卡
net.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型,可指定多张显卡

只指定主显卡,如下:

import torch
import os
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

torch.cuda常用指令

指定特定显卡,如下:

import torch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
x = torch.tensor([[1,2,3],[4,5,6]]).cuda('cuda:1')
print(x.device)

torch.cuda常用指令