使用TensorBoard可视化YOLOP模型结构失败的问题求助

错误原因 #

你遇到的List inputs to traced functions must have consistent element type. Found Tuple[Tensor, List[Tensor]] and Tensor错误，是因为YOLOP模型的forward方法返回结果中包含List[Tensor]类型（比如检测分支的预测框列表），而SummaryWriter.add_graph底层依赖的torch.jit.trace要求追踪的函数输出必须是张量或张量的统一容器类型，混合的列表与元组/张量会导致类型不匹配，无法完成模型结构追踪。

解决方案 #

方法1：创建模型包装类，统一输出为张量类型 #

定义一个包装类，调整原模型的输出结构，只返回可被追踪的张量（或张量元组）：

import torch

class YOLOPWrapper(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model
    
    def forward(self, x):
        # 原YOLOP通常返回(detection_output, segmentation_output, lane_output)
        det_out, seg_out, lane_out = self.model(x)
        # 若det_out是List[Tensor]，可选择提取其中的张量元素，或仅返回分割/车道线分支的输出
        # 这里示例仅返回分割和车道线分支的张量输出，用于结构可视化
        return seg_out, lane_out

修改可视化代码，使用包装后的模型：

if __name__ == "__main__":
    from torch.utils.tensorboard import SummaryWriter
    import torch
    model = get_net(False)
    wrapped_model = YOLOPWrapper(model)
    input_ = torch.randn((1, 3, 256, 256))
    with SummaryWriter(comment="YOLOP") as w:
        w.add_graph(wrapped_model, (input_,))

方法2：关闭严格追踪模式（简易但可能丢失部分结构） #

在add_graph中添加strict=False参数，允许追踪器忽略部分动态结构：

if __name__ == "__main__":
    from torch.utils.tensorboard import SummaryWriter
    model = get_net(False)
    input_ = torch.randn((1, 3, 256, 256))
    with SummaryWriter(comment="YOLOP") as w:
        w.add_graph(model, (input_,), strict=False)

说明 #

• 方法1更可靠，能完整追踪模型的核心结构，适合需要准确可视化的场景；
• 方法2操作简单，但可能无法追踪模型中的动态分支（如基于条件判断的结构），仅作为临时解决方案。

内容的提问来源于stack exchange，提问作者Kerry killer