pytorch는 텐서와 모델이 어떤 GPU에 있는지 확인하고 GPU 장치를 지정합니다.

1. 텐서가 있는 장치 보기:

data = data.cuda()#将数据转移到gpu上
 
print(data.device)  # 输出:cuda:0
 
data = data.cpu()#将数据转移到cpu上
 
print(data.device)  # 输出:cpu

2. 모델이 위치한 기기 보기

model = model.cuda()#将模型转移到gpu上
 
print(next(model.parameters()).device)  # 输出:cuda:0
 
model = model.cpu()#将模型转移到cpu上
 
print(next(model.parameters()).device)  # 输出:cpu

3. Pytorch에서 모델과 텐서를 GPU에 로드하는 두 가지 일반적인 방법이 있습니다.

방법 1:

# 如果GPU可用,将模型和张量加载到GPU上
if torch.cuda.is_available():
    model = model.cuda()
    x = x.cuda()
    y = y.cuda()

방법 2:

# 分配到的GPU或CPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将模型加到GPU
model=model.to(device)
# 将张量加到GPU
x=x.to(device)
y=y.to(device)

4. GPU 코드 지정

# 代码1:
torch.cuda.set_device(1)

# 代码2:
device = torch.device("cuda:1")

# 代码3:(官方推荐使用),
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

(如果你想同时调用两块GPU的话)
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'

참조 링크: PyTorch에서 지정된 GPU 선택

지정된 GPU 코드는 다음 그림과 같이 프로그램 세그먼트의 시작 부분에 배치되어야 합니다.
gpu

5. GPU 수 보기

torch.cuda.device_count()

추천

출처blog.csdn.net/flyingluohaipeng/article/details/128431003