在自己的代码上使用 Tensorboard 可视化，处理多张输入图像时遇到的问题 RuntimeError: 0 INTERNAL ASSERT FAILED at ..\torch\csrc\jit

问题描述

在自己的代码上使用 Tensorboard 进行训练可视化时（具体过程参考教程 https://blog.csdn.net/prague6695/article/details/119738030），发生如下错误：RuntimeError: 0 INTERNAL ASSERT FAILED at ..\torch\csrc\jit\ir\alias_analysis.cpp:318, please report a bug to PyTorch. We don't have an op for aten::to but it isn't a special case. Argument types: Tensor, None, int, Device, bool, bool, bool, int,

问题分析

由于自己的网络有 4 个图片输入，所以有如下这段代码，与参考教程不一致：

    # Tensorboard 使用步骤 2，实将模型写入 tensorboard 中
    init_img = torch.zeros((1, 3, 112, 112), device=device)  # 就是指定输入图片的大小。
    face_img = torch.zeros((1, 3, 224, 224), device=device)
    rect_img = torch.zeros((1, 12), device=device)
    tb_writer.add_graph(net, (init_img, init_img, face_img, rect_img))

运行训练文件后报错信息如开头提到的。

问题解决

将最后一行代码修改为：

tb_writer.add_graph(net.module, (init_img, init_img, face_img, rect_img))

参考链接：https://github.com/pytorch/pytorch/issues/35180