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

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

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

你們可能遇到這樣子的困擾:好比說運行本身編寫的 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 語句才能找到緣由。oop

TorchSnooper 就是一個設計了用來解決這個問題的工具。TorchSnooper 的安裝很是簡單,只須要執行標準的 Python 包安裝指令就好:測試

pip install torchsnooper

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

安裝完了之後,下面就用兩個例子來講明一下怎麼使用。ui

例子1scala

好比說咱們寫了一個很是簡單的函數:設計

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 了。大功告成。

相關文章
相關標籤/搜索