在PyTorch中,可以使用torch.cuda.set_device()
方法來指定GPU訓練環境變量。以下是具體步驟:
torch
和torch.cuda
模塊:import torch
import torch.cuda
torch.cuda.set_device()
方法指定使用的GPU設備。例如,如果要使用第一個GPU設備,可以這樣設置:torch.cuda.set_device(0)
torch.cuda.set_device()
方法,確保模型和數據都在指定的GPU上進行訓練。torch.cuda.set_device(0)
model = Model()
這樣就可以在PyTorch中指定GPU訓練環境變量了。需要注意的是,如果要在多個GPU上進行訓練,可以使用torch.nn.DataParallel
或torch.nn.parallel.DistributedDataParallel
來實現多GPU并行訓練。