關於pytorch的調試工具之torchsnooper

1. 代碼

import torch
def my_fun(mask, x):
    y = torch.zeros(5, 6)
    y.masked_fill_(mask, x)
    return y

x = [5, 4, 3, 2, 1]
mask = torch.zeros(5, 6, dtype=torch.float)
for e_id, src_len in enumerate(x):
    mask[e_id, src_len:] = 1
mask = mask.to(device='cuda')
x = torch.tensor(1, device='cuda')
y = my_fun(mask, x)
print(y)

2. 運行出問題

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

點開拋異常的地方爲y.masked_fill_(mask,x)這行代碼出錯了，估計是哪個變量沒有放到CUDA，爲了查得仔細，增加torchsnooper查看一下。

3. 增加工具

import torch
import torchsnooper

@torchsnooper.snoop()
def my_fun(mask, x):
    y = torch.zeros(5, 6)
    y.masked_fill_(mask, x)
    return y

x = [5, 4, 3, 2, 1]
mask = torch.zeros(5, 6, dtype=torch.float)
for e_id, src_len in enumerate(x):
    mask[e_id, src_len:] = 1
mask = mask.to(device='cuda')
x = torch.tensor(1, device='cuda')
y = my_fun(mask, x)
print(y)

這個函數輸出內容爲：

Starting var:.. mask = tensor<(5, 6), float32, cuda:0>
Starting var:.. x = tensor<(), int64, cuda:0>
02:40:01.684679 call        23 def my_fun(mask, x):
02:40:01.691078 line        24     y = torch.zeros(5, 6)
New var:....... y = tensor<(5, 6), float32, cpu>
02:40:01.691401 line        27     y.masked_fill_(mask, x)
02:40:01.694904 exception   27     y.masked_fill_(mask, x)
RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 'mask'

從內容中可以看到代碼的運行過程，很快找到了y值是保存在cpu上，加上

y = y.to(device='cuda')

再運行,又出現了一個問題：

Starting var:.. mask = tensor<(5, 6), float32, cuda:0>
Starting var:.. x = tensor<(), int64, cuda:0>
02:42:15.633153 call        23 def my_fun(mask, x):
02:42:15.641048 line        24     y = torch.zeros(5, 6)
New var:....... y = tensor<(5, 6), float32, cpu>
02:42:15.641678 line        25     y = y.to(device='cuda')
Modified var:.. y = tensor<(5, 6), float32, cuda:0>
02:42:15.642666 line        27     y.masked_fill_(mask, x)
02:42:15.651113 exception   27     y.masked_fill_(mask, x)
RuntimeError: Expected object of scalar type Byte but got scalar type Float for argument #2 'mask'

這裏看到一個信息，是類型不對；mask的數據類型不對，masked_fill_函數的masked參數類型爲ByteTensor，故得把mask轉一下類型mask.byte()或者在mask創建時設置mask = torch.zeros(5, 6, dtype=torch.uint8)；修改完沒有問題了。

附1

Data tyoe	CPU tensor	GPU tensor
32-bit floating point	torch.FloatTensor	torch.cuda.FloatTensor
64-bit floating point	torch.DoubleTensor	torch.cuda.DoubleTensor
16-bit floating point	N/A	torch.cuda.HalfTensor
8-bit integer (unsigned)	torch.ByteTensor	torch.cuda.ByteTensor
8-bit integer (signed)	torch.CharTensor	torch.cuda.CharTensor
16-bit integer (signed)	torch.ShortTensor	torch.cuda.ShortTensor
32-bit integer (signed)	torch.IntTensor	torch.cuda.IntTensor
64-bit integer (signed)	torch.LongTensor	torch.cuda.LongTensor

附2

masked_fill_(mask, value)
在mask值爲1的位置處用value填充。mask的元素個數需和本tensor相同，但尺寸可以不同。

參數： - mask (ByteTensor)-二進制掩碼 - value (Tensor)-用來填充的值

附3

TorchSnooper源碼與安裝使用詳情：
https://github.com/zasdfgbnm/TorchSnooper