PyTorch程式码除错利器：自动print每行程式码的Tensor资讯
机器之心释出
作者：zasdfgbnm
本文介绍一个用于 PyTorch 程式码的实用工具 TorchSnooper。作者是TorchSnooper的作者，也是PyTorch开发者之一
GitHub 专案地址： https://github.com/zasdfgbnm/TorchSnooper
大家可能遇到这样子的困扰：比如说执行自己编写的 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 语句才能找到原因。
TorchSnooper 就是一个设计了用来解决这个问题的工具。TorchSnooper 的安装非常简单，只需要执行标准的 Python 包安装指令就好：
pip install torchsnooper
安装完了以后，只需要用 @torchsnooper.snoop() 装饰一下要除错的函式，这个函式在执行的时候，就会自动 print 出来每一行的执行结果的 tensor 的形状、资料型别、装置、是否需要梯度的资讯。
安装完了以后，下面就用两个例子来说明一下怎么使用。
例子1
比如说我们写了一个非常简单的函式：
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 了。大功告成。
