关于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