pytorch怎么自動(dòng)打印每行代碼的tensor信息？

我們?cè)谶M(jìn)行pytorch調(diào)試的時(shí)候可能會(huì)遇到一些錯(cuò)誤，當(dāng)遇到這些錯(cuò)誤的時(shí)候我們需要去打印tensor信息，但是因?yàn)槲覀儾恢绬栴}出在哪，所以我們可能需要打印很多很多的tensor信息，手寫這么多tensor信息的print語(yǔ)句是相當(dāng)累人（且低效）的，今天小編帶來一個(gè)pytorch調(diào)試工具——tensorsnooper，這樣就能實(shí)現(xiàn)pytorch自動(dòng)打印每行代碼的tensor信息了。

本文介紹一個(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 語(yǔ)句，非常麻煩。

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

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

pip install torchsnooper

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

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

例子1

比如說我們寫了一個(gè)非常簡(jiǎn)單的函數(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 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)

然后運(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è)變量。這一行對(duì)應(yīng)代碼中的 y = torch.zeros(6)。于是我們意識(shí)到，在使用 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è)簡(jiǎn)單的線性模型：

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

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

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

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

運(yùn)行程序，我們看到了一長(zhǎng)串的輸出，一點(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:....... _ = 002: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) ** 2New 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.97929000854492202: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ù)起來也很簡(jiǎn)單，把 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:....... _ = 002: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) ** 2New 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)過測(cè)試，loss 現(xiàn)在已經(jīng)可以降到很接近 0 了。大功告成。

以上就是pytorch自動(dòng)打印每行代碼的tensor信息的方法了，希望能給大家一個(gè)參考，也希望大家多多支持W3Cschool。