pytorch中画模型图

需要导入tensorboard包
pip install tensoboard

代码

from torch.utils.tensorboard import SummaryWriter
model = FPN()

with SummaryWriter(comment="FPN") as w:
    w.add_graph(model=model, input_to_model=torch.rand(15, 3, 224, 224))

运行之后开一个终端，输入tensorboard --logdir=runs，再按照它提示的地址打开浏览器