Google colab 上使用TPU进行多核训练时该怎么调试（Pytorch）

问题

 

使用Google colab并不像使用pycharm那么容易调试，特别是文件内容比较多且杂的时候，调试起来极为不方便。而使用Google云TPU的八个核进行训练更是使得调试工作难上加难，到出错的地方要么直接退出，要么就卡住不动，错误信息都不给一个，让人很是火大。

关于调试

我并没有直接在colab上运行代码，而是把python文件上传到Google云盘上，然后再在colab上挂载云盘，最后运行python文件。如：

!python3 test.py

所以我的调试代码是写在python文件中的。

# 导入包
from IPython.core.debugger import set_trace

for i in range(10):
    set_trace()  # 设置断点
    i += 1

运行python文件，程序会在断点处停下。输入如下命令可以进行调试：

按键 (可以全名， 也可以缩写)	功能
ENTER	重复上次命令
c (cont) (continue)	继续程序
l (list)	可以列出当前将要运行的代码块
s (step)	步入
n (next)	步过
r (return)	运行到程序结束
!<python 命令>	执行命令
h (help)	帮助
h + 指令	显示指令帮助内容
a (args)	打印当前函数的参数
j + 行数	跳转到指定的行数
p (print)	打印变量
q (quit)	退出调试

使用云TPU怎么调试 

问题：使用多核训练就像使用了多线程一样，每个进程干着自己的事情，你就是一个外人，根本无法插入进去。

使用教程：（需要科学上网）

PyTorch on Cloud TPUs: Single Core Training AlexNet on Fashion MNIST

PyTorch on Cloud TPUs: MultiCore Training AlexNet on Fashion MNIST

解决方法：其实解决的办法很简单，就是使用单核进行调试，之后再改为多核运行。断点插到map_fn()函数中。

# 原来的
# xmp.spawn(map_fn, args=(flags,), nprocs=8, start_method='fork')

xmp.spawn(map_fn, args=(flags,), nprocs=1, start_method='fork')

在运行过程中出现的问题 

1、运行loss.item()的时候出现阻塞，这时程序不会报错，而是会一直堵着。

原因：loss的类型为

tensor(16.1250, device='xla:1', dtype=torch.bfloat16, grad_fn=<DivBackward0>)

 所以应改为loss.cpu().item()

2、主进程单独执行程序时会堵塞，因为不能在不同的进程上执行不同的计算，而在cpu上进行运算是可以的。

# Don't perform different computations on different processes!
# Warning: uncommenting the below and running this cell will likely hang your Colab!
def simple_map_fn(index, flags):
  torch.manual_seed(1234)
  device = xm.xla_device()  

  if xm.is_master_ordinal():
    t = torch.randn((2, 2), device=device)  # Divergent Cloud TPU computation!


xmp.spawn(simple_map_fn, args=(flags,), nprocs=8, start_method='fork')

 解决办法：使用do_on_ordinals()函数，这个函数会把data转换为cpu数据进行运算。

需要注意的问题

使用ipdb调试不会像pycharm那样可以随便进出，有时候执行命令s根本不会步入，而是直接执行了。而且在本应该卡住的的地方还能继续执行。所以有时候最好配合print使用。