PyTorch代碼調試利器：自動print每行代碼的Tensor信息

機器之心發布

作者：zasdfgbnm

本文介紹一個用於 PyTorch 代碼的實用工具 TorchSnooper。作者是TorchSnooper的作者，也是PyTorch開發者之一

GitHub 項目地址： https://github.com/zasdfgbnm/TorchSnooper

大家可能遇到這樣子的困擾：比如說運行自己編寫的 PyTorch 代碼的時候，PyTorch 提示你說數據類型不匹配，需要一個 double 的 tensor 但是你給的卻是 float；再或者就是需要一個 CUDA tensor, 你給的卻是個 CPU tensor。比如下面這種：

RuntimeError: Expected object of scalar type Double but got scalar type Float

這種問題調試起來很麻煩，因為你不知道從哪裡開始出問題的。比如你可能在代碼的第三行用 torch.zeros 新建了一個 CPU tensor, 然後這個 tensor 進行了若干運算，全是在 CPU 上進行的，一直沒有報錯，直到第十行需要跟你作為輸入傳進來的 CUDA tensor 進行運算的時候，才報錯。要調試這種錯誤，有時候就不得不一行行地手寫 print 語句，非常麻煩。

再或者，你可能腦子裡想像著將一個 tensor 進行什麼樣子的操作，就會得到什麼樣子的結果，但是 PyTorch 中途報錯說 tensor 的形狀不匹配，或者壓根沒報錯但是最終出來的形狀不是我們想要的。這個時候，我們往往也不知道是什麼地方開始跟我們「預期的發生偏離的」。我們有時候也得需要插入一大堆 print 語句才能找到原因。

TorchSnooper 就是一個設計了用來解決這個問題的工具。TorchSnooper 的安裝非常簡單，只需要執行標準的 Python 包安裝指令就好：

pip install torchsnooper

安裝完了以後，只需要用 @torchsnooper.snoop() 裝飾一下要調試的函數，這個函數在執行的時候，就會自動 print 出來每一行的執行結果的 tensor 的形狀、數據類型、設備、是否需要梯度的信息。

安裝完了以後，下面就用兩個例子來說明一下怎麼使用。

例子1

比如說我們寫了一個非常簡單的函數：

def myfunc(mask, x):
y = torch.zeros(6)
y.masked_scatter_(mask, x)
return y

我們是這樣子使用這個函數的：

mask = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda")
source = torch.tensor([1.0, 2.0, 3.0], device="cuda")
y = myfunc(mask, source)

上面的代碼看起來似乎沒啥問題，然而實際上跑起來，卻報錯了：

RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 "mask"

問題在哪裡呢？讓我們 snoop 一下！用 @torchsnooper.snoop() 裝飾一下 myfunc 函數：

import torch
import torchsnooper
@torchsnooper.snoop()
def myfunc(mask, x):
y = torch.zeros(6)
y.masked_scatter_(mask, x)
return y
mask = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda")
source = torch.tensor([1.0, 2.0, 3.0], device="cuda")
y = myfunc(mask, source)

然後運行我們的腳本，我們看到了這樣的輸出：

Starting var:.. mask = tensor<(6,), int64, cuda:0>
Starting var:.. x = tensor<(3,), float32, cuda:0>
21:41:42.941668 call 5 def myfunc(mask, x):
21:41:42.941834 line 6 y = torch.zeros(6)
New var:....... y = tensor<(6,), float32, cpu>
21:41:42.943443 line 7 y.masked_scatter_(mask, x)
21:41:42.944404 exception 7 y.masked_scatter_(mask, x)

結合我們的錯誤，我們主要去看輸出的每個變數的設備，找找最早從哪個變數開始是在 CPU 上的。我們注意到這一行：

New var:....... y = tensor<(6,), float32, cpu>

這一行直接告訴我們，我們創建了一個新變數 y, 並把一個 CPU tensor 賦值給了這個變數。這一行對應代碼中的 y = torch.zeros(6)。於是我們意識到，在使用 torch.zeros 的時候，如果不人為指定設備的話，默認創建的 tensor 是在 CPU 上的。我們把這一行改成 y = torch.zeros(6, device="cuda")，這一行的問題就修復了。

這一行的問題雖然修復了，我們的問題並沒有解決完整，再跑修改過的代碼還是報錯，但是這個時候錯誤變成了：

RuntimeError: Expected object of scalar type Byte but got scalar type Long for argument #2 "mask"

好吧，這次錯誤出在了數據類型上。這次錯誤報告比較有提示性，我們大概能知道是我們的 mask 的數據類型錯了。再看一遍 TorchSnooper 的輸出，我們注意到：

Starting var:.. mask = tensor<(6,), int64, cuda:0>

果然，我們的 mask 的類型是 int64, 而不應該是應有的 uint8。我們把 mask 的定義修改好：

mask = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda", dtype=torch.uint8)

然後就可以運行了。

例子 2

這次我們要構建一個簡單的線性模型：

model = torch.nn.Linear(2, 1)

我們想要擬合一個平面 y = x1 + 2 * x2 + 3，於是我們創建了這樣一個數據集：

x = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
y = torch.tensor([3.0, 5.0, 4.0, 6.0])

我們使用最普通的 SGD 優化器來進行優化，完整的代碼如下：

import torch
model = torch.nn.Linear(2, 1)
x = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
y = torch.tensor([3.0, 5.0, 4.0, 6.0])
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for _ in range(10):
optimizer.zero_grad()
pred = model(x)
squared_diff = (y - pred) ** 2
loss = squared_diff.mean()
print(loss.item())
loss.backward()
optimizer.step()

然而運行的過程我們發現，loss 降到 1.5 左右就不再降了。這是很不正常的，因為我們構建的數據都是無誤差落在要擬合的平面上的，loss 應該降到 0 才算正常。

乍看上去，不知道問題在哪裡。抱著試試看的想法，我們來 snoop 一下子。這個例子中，我們沒有自定義函數，但是我們可以使用 with 語句來激活 TorchSnooper。把訓練的那個循環裝進 with 語句中去，代碼就變成了：

import torch
import torchsnooper
model = torch.nn.Linear(2, 1)
x = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
y = torch.tensor([3.0, 5.0, 4.0, 6.0])
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
with torchsnooper.snoop():
for _ in range(10):
optimizer.zero_grad()
pred = model(x)
squared_diff = (y - pred) ** 2
loss = squared_diff.mean()
print(loss.item())
loss.backward()
optimizer.step()

運行程序，我們看到了一長串的輸出，一點一點瀏覽，我們注意到

New var:....... model = Linear(in_features=2, out_features=1, bias=True)
New var:....... x = tensor<(4, 2), float32, cpu>
New var:....... y = tensor<(4,), float32, cpu>
New var:....... optimizer = SGD (Parameter Group 0 dampening: 0 lr: 0....omentum: 0 nesterov: False weight_decay: 0)
02:38:02.016826 line 12 for _ in range(10):
New var:....... _ = 0
02:38:02.017025 line 13 optimizer.zero_grad()
02:38:02.017156 line 14 pred = model(x)
New var:....... pred = tensor<(4, 1), float32, cpu, grad>
02:38:02.018100 line 15 squared_diff = (y - pred) ** 2
New var:....... squared_diff = tensor<(4, 4), float32, cpu, grad>
02:38:02.018397 line 16 loss = squared_diff.mean()
New var:....... loss = tensor<(), float32, cpu, grad>
02:38:02.018674 line 17 print(loss.item())
02:38:02.018852 line 18 loss.backward()
26.979290008544922
02:38:02.057349 line 19 optimizer.step()

仔細觀察這裡面各個 tensor 的形狀，我們不難發現，y 的形狀是 (4,)，而 pred 的形狀卻是 (4, 1)，他們倆相減，由於廣播的存在，我們得到的 squared_diff 的形狀就變成了 (4, 4)。

這自然不是我們想要的結果。這個問題修復起來也很簡單，把 pred 的定義改成 pred = model(x).squeeze() 即可。現在再看修改後的代碼的 TorchSnooper 的輸出：

New var:....... model = Linear(in_features=2, out_features=1, bias=True)
New var:....... x = tensor<(4, 2), float32, cpu>
New var:....... y = tensor<(4,), float32, cpu>
New var:....... optimizer = SGD (Parameter Group 0 dampening: 0 lr: 0....omentum: 0 nesterov: False weight_decay: 0)
02:46:23.545042 line 12 for _ in range(10):
New var:....... _ = 0
02:46:23.545285 line 13 optimizer.zero_grad()
02:46:23.545421 line 14 pred = model(x).squeeze()
New var:....... pred = tensor<(4,), float32, cpu, grad>
02:46:23.546362 line 15 squared_diff = (y - pred) ** 2
New var:....... squared_diff = tensor<(4,), float32, cpu, grad>
02:46:23.546645 line 16 loss = squared_diff.mean()
New var:....... loss = tensor<(), float32, cpu, grad>
02:46:23.546939 line 17 print(loss.item())
02:46:23.547133 line 18 loss.backward()
02:46:23.591090 line 19 optimizer.step()

現在這個結果看起來就正常了。並且經過測試，loss 現在已經可以降到很接近 0 了。大功告成。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！