对于模型的输出是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类型)。参考如下: