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
,再按照它提示的地址打开浏览器