pytorch中tensorboard网络可视化输出错误:

        对于模型的输出是list或者dict,需要使用可视化展示model时会产生错误,仅支持tuple或者namedtuple输出方式:

RuntimeError: Encountering a dict at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.

 可以通过如下方法将模型的输出进行更改:

from collections import namedtuple
from typing import Any

# pylint: disable = abstract-method
class ModelWrapper(torch.nn.Module):
    """
    Wrapper class for model with dict/list rvalues.
    """

    def __init__(self, model: torch.nn.Module) -> None:
        """
        Init call.
        """
        super().__init__()
        self.model = model

    def forward(self, input_x: torch.Tensor) -> Any:
        """
        Wrap forward call.
        """
        data = self.model(input_x)

        if isinstance(data, dict):
            data_named_tuple = namedtuple("ModelEndpoints", sorted(data.keys()))  # type: ignore
            data = data_named_tuple(**data)  # type: ignore

        elif isinstance(data, list):
            data = tuple(data)

        return data

在修改网络的测试输入input,并将你自己的网络模型的类输入到此类中: 

#以下内容是我自己网络测试使用
if __name__ == '__main__':
    writer = SummaryWriter('logs')
    input = torch.ones((1,3,800,320))
    net = parsingNet(pretrained=True, backbone='50',num_grid_row=56,num_cls_row =100, num_grid_col = 40, num_cls_col = 100, num_lane_on_row = 2, num_lane_on_col = 2,input_height = 800, input_width = 320)
    model_wrapper = ModelWrapper(net)
    writer.add_graph(model_wrapper, input)
    # output = net(input)
    # print(type(output))
    # print(output.keys())
    # print(isinstance(output,dict))
    
    # writer.add_graph(net,input)
    writer.close()

 通过以上修改就可以可视化输出网络(我的网络输出的是dict类型)。参考如下:

  • 2
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值