PyTorch程式碼除錯利器: 自動print每行程式碼的Tensor資訊

（點選上方快速關註並設定為星標，一起學Python）

來源：機器之心  作者：zasdfgbnm   連結：

https://mp.weixin.qq.com/s/PVIWWIqbZuEe4lwVMDYA5Q

 

本文介紹一個用於 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 = tensor6,), int64, cuda:0>
Starting var:.. x = tensor3,), 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 = tensor6,), 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 = tensor6,), 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 = tensor6,), 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 = tensor4, 2), float32, cpu>
New var:....... y = tensor4,), 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 = tensor4, 1), float32, cpu, grad>
02:38:02.018100 line        15         squared_diff = (y - pred) ** 2
New var:....... squared_diff = tensor4, 4), float32, cpu, grad>
02:38:02.018397 line        16         loss = squared_diff.mean()
New var:....... loss = tensor
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 = tensor4, 2), float32, cpu>
New var:....... y = tensor4,), 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 = tensor4,), float32, cpu, grad>
02:46:23.546362 line        15         squared_diff = (y - pred) ** 2
New var:....... squared_diff = tensor4,), float32, cpu, grad>
02:46:23.546645 line        16         loss = squared_diff.mean()
New var:....... loss = tensor
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 了。大功告成。