本文介紹一個用于 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?ymask?=?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?torchmodel?=?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)?**?2loss?=?squared_diff.mean()print(loss.item())loss.backward()optimizer.step()

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

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

import?torch
import?torchsnoopermodel?=?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)?**?2loss?=?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 了。大功告成。

總結

以上是生活随笔為你收集整理的PyTorch代码调试利器: 自动print每行代码的Tensor信息的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。