本文介紹一個(gè)用于 PyTorch 代碼的實(shí)用工具 TorchSnooper。作者是TorchSnooper的作者，也是PyTorch開發(fā)者之一。

GitHub 項(xiàng)目地址： https://github.com/zasdfgbnm/TorchSnooper

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

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

這種問題調(diào)試起來很麻煩，因?yàn)槟悴恢缽哪睦镩_始出問題的。比如你可能在代碼的第三行用 torch.zeros 新建了一個(gè) CPU tensor, 然后這個(gè) tensor 進(jìn)行了若干運(yùn)算，全是在 CPU 上進(jìn)行的，一直沒有報(bào)錯(cuò)，直到第十行需要跟你作為輸入傳進(jìn)來的 CUDA tensor 進(jìn)行運(yùn)算的時(shí)候，才報(bào)錯(cuò)。要調(diào)試這種錯(cuò)誤，有時(shí)候就不得不一行行地手寫 print 語句，非常麻煩。

再或者，你可能腦子里想象著將一個(gè) tensor 進(jìn)行什么樣子的操作，就會得到什么樣子的結(jié)果，但是 PyTorch 中途報(bào)錯(cuò)說 tensor 的形狀不匹配，或者壓根沒報(bào)錯(cuò)但是最終出來的形狀不是我們想要的。這個(gè)時(shí)候，我們往往也不知道是什么地方開始跟我們「預(yù)期的發(fā)生偏離的」。我們有時(shí)候也得需要插入一大堆 print 語句才能找到原因。

TorchSnooper 就是一個(gè)設(shè)計(jì)了用來解決這個(gè)問題的工具。TorchSnooper 的安裝非常簡單，只需要執(zhí)行標(biāo)準(zhǔn)的 Python 包安裝指令就好：

pip?install?torchsnooper

安裝完了以后，只需要用 @torchsnooper.snoop() 裝飾一下要調(diào)試的函數(shù)，這個(gè)函數(shù)在執(zhí)行的時(shí)候，就會自動(dòng) print 出來每一行的執(zhí)行結(jié)果的 tensor 的形狀、數(shù)據(jù)類型、設(shè)備、是否需要梯度的信息。

安裝完了以后，下面就用兩個(gè)例子來說明一下怎么使用。

例子1

比如說我們寫了一個(gè)非常簡單的函數(shù)：

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

我們是這樣子使用這個(gè)函數(shù)的：

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)

上面的代碼看起來似乎沒啥問題，然而實(shí)際上跑起來，卻報(bào)錯(cuò)了：

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

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

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)

然后運(yùn)行我們的腳本，我們看到了這樣的輸出：

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)

結(jié)合我們的錯(cuò)誤，我們主要去看輸出的每個(gè)變量的設(shè)備，找找最早從哪個(gè)變量開始是在 CPU 上的。我們注意到這一行：

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

這一行直接告訴我們，我們創(chuàng)建了一個(gè)新變量 y, 并把一個(gè) CPU tensor 賦值給了這個(gè)變量。這一行對應(yīng)代碼中的 y = torch.zeros(6)。于是我們意識到，在使用 torch.zeros 的時(shí)候，如果不人為指定設(shè)備的話，默認(rèn)創(chuàng)建的 tensor 是在 CPU 上的。我們把這一行改成 y = torch.zeros(6, device='cuda')，這一行的問題就修復(fù)了。

這一行的問題雖然修復(fù)了，我們的問題并沒有解決完整，再跑修改過的代碼還是報(bào)錯(cuò)，但是這個(gè)時(shí)候錯(cuò)誤變成了：

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

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

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

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

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

然后就可以運(yùn)行了。

例子 2

這次我們要構(gòu)建一個(gè)簡單的線性模型：

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

我們想要擬合一個(gè)平面 y = x1 + 2 * x2 + 3，于是我們創(chuàng)建了這樣一個(gè)數(shù)據(jù)集：

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 優(yōu)化器來進(jìn)行優(yōu)化，完整的代碼如下：

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()

然而運(yùn)行的過程我們發(fā)現(xiàn)，loss 降到 1.5 左右就不再降了。這是很不正常的，因?yàn)槲覀儤?gòu)建的數(shù)據(jù)都是無誤差落在要擬合的平面上的，loss 應(yīng)該降到 0 才算正常。

乍看上去，不知道問題在哪里。抱著試試看的想法，我們來 snoop 一下子。這個(gè)例子中，我們沒有自定義函數(shù)，但是我們可以使用 with 語句來激活 TorchSnooper。把訓(xùn)練的那個(gè)循環(huán)裝進(jìn) 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()

運(yùn)行程序，我們看到了一長串的輸出，一點(diǎn)一點(diǎn)瀏覽，我們注意到

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()

仔細(xì)觀察這里面各個(gè) tensor 的形狀，我們不難發(fā)現(xiàn)，y 的形狀是 (4,)，而 pred 的形狀卻是 (4, 1)，他們倆相減，由于廣播的存在，我們得到的 squared_diff 的形狀就變成了 (4, 4)。

這自然不是我們想要的結(jié)果。這個(gè)問題修復(fù)起來也很簡單，把 pred 的定義改成 pred = model(x).squeeze() 即可。現(xiàn)在再看修改后的代碼的 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()

現(xiàn)在這個(gè)結(jié)果看起來就正常了。并且經(jīng)過測試，loss 現(xiàn)在已經(jīng)可以降到很接近 0 了。大功告成。

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。