选择服务器上的第二个GPU

Question

我在服务器上运行代码。那里有2个gpu，第一个是忙的。然而，我找不到在它们之间切换的方法。我正在使用pytorch，如果这是重要的。以下几行代码应该修改:

device = 'cuda' if torch.cuda.is_available() else 'cpu'

修改只能在这里说明。谢谢。

Answer 1

cuda默认选择cuda:0，可通过cuda:1切换到其他GPU

所以，你的行变成:

device = 'cuda:1' if torch.cuda.is_available() else 'cpu'

你可以阅读更多关于CUDA语义的信息。

Answer 2

这是我使用FastAI和预训练模型进行推理时的方法。首先，当使用fai (import fastai.vision.all as fai)定义模型时，我获得模型实例并将其放入指定的GPU(例如使用gpu_id=3):

model = fai.nn.Sequential(body, head)
model.cuda(device=gpu_id)

然后，当加载模型权重时，我也指定要使用哪个设备(否则它会在GPU 0中创建模型的副本):

model.load_state_dict(torch.load(your_model_state_filepath, torch.device(gpu_id)))