[PyTorch] 判斷當前資料所使用的 GPU

Last Updated on 2021-07-05 by Clay

在我使用 PyTorch 訓練模型的時候，經常會發生我使用 GPU_A 去訓練模型、儲存模型，然而在測試模型效果的時候，卻不小心使用到了 GPU_B 來讀取測試資料 (我有多片 GPU 可以使用，還滿奢侈的 XDD)，然後再用已經儲存好的模型 (GPU_A 訓練) 來測試 ——

這樣的結果會是，你拿清朝的劍斬明朝的官 …… 使用 GPU_A 上的模型去計算 GPU_B 上的資料。然後，我們會吃下一個 Error Message，顯示兩邊資料不在同樣的設備上。

當然，這是個可以解決的問題，比方說把要跑的資料全部都固定丟到某個裝置底下。比方我我全部丟到第一片 GPU (cuda:0) 好了。

model = model.to('cuda:0')

model = model.to(‘cuda:0’)

不過我想知道的是，究竟有沒有辦法直接看我的資料位於哪個裝置上呢？上網找了一下，發現還真的有人有這個需求，故順手紀錄在這裡。

使用 get_device() 查看

註：此方法只對 Tensor 有用，而且對於還在 CPU 上的 Tensor 好像不起作用的樣子。

import torch

a = torch.tensor([5, 3]).to('cuda:3')
print(a.get_device())

import torch

a = torch.tensor([5, 3]).to(‘cuda:3’)
print(a.get_device())

Output:

可以確認這方法是可行的。