使用os.environ["CUDA_VISIBLE_DEVICES"]

這種方法是通過環(huán)境變量限制可見的CUDA設備，從而在多個GPU的機器上只讓PyTorch看到并使用指定的GPU。這種方式的好處是所有后續(xù)的CUDA調用都會使用這個GPU，并且代碼中不需要顯式地指定設備索引。

import os# 設置只使用2號GPU
os.environ["CUDA_VISIBLE_DEVICES"] = '2'import torch
import torch.nn as nn# 檢查PyTorch是否檢測到GPU
if torch.cuda.is_available():print(f"Using GPU: {torch.cuda.get_device_name(0)}")  # 注意這里是0，因為只有一個可見的GPU
else:print("No GPU available, using CPU instead.")# 定義模型
class YourModel(nn.Module):def __init__(self):super(YourModel, self).__init__()self.layer = nn.Linear(10, 1)def forward(self, x):return self.layer(x)# 創(chuàng)建模型并移動到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = YourModel().to(device)# 示例數據和前向傳播
input_data = torch.randn(5, 10).to(device)
output = model(input_data)
print(output)

直接指定設備索引

這種方法是在代碼中直接指定要使用的設備索引，無需修改環(huán)境變量。這種方式更加顯式，并且可以在同一腳本中使用多個不同的GPU。

import torch
import torch.nn as nn# 檢查設備是否可用并打印設備名稱
if torch.cuda.is_available():device = torch.device("cuda:2")  # 直接指定設備索引print(f"Using GPU: {torch.cuda.get_device_name(2)}")
else:device = torch.device("cpu")print("No GPU available, using CPU instead.")# 定義模型
class YourModel(nn.Module):def __init__(self):super(YourModel, self).__init__()self.layer = nn.Linear(10, 1)def forward(self, x):return self.layer(x)# 創(chuàng)建模型并移動到指定的GPU
model = YourModel().to(device)# 示例數據和前向傳播
input_data = torch.randn(5, 10).to(device)
output = model(input_data)
print(output)