Tensorboard 可视化模型计算图

Posted 2022-11-23 洪流之源

通过tensorboard的add_graph接口可实现模型计算图的可视化：

from torch.utils.tensorboard import SummaryWriter
import torch
import torchvision

if __name__ == '__main__':

    summary_writer = SummaryWriter(log_dir='log_graph', comment='test tensorboard graph', filename_suffix='test_1')

    image = torch.randn(1, 3, 32, 32)

    model = torchvision.models.vgg16(num_classes=2)

    # 可视化模型计算图
    # model : 模型，必须是nn.Module
    # input_to_model : 输出给模型的数据
    # verbose : 是否打印计算图结构信息
    summary_writer.add_graph(model=model, input_to_model=image, verbose=False)

    summary_writer.close()

运行tensorboard，并指定日志目录：

tensorboard --logdir log_graph

如下是计算图的可视化结果，双击计算图中的节点可查看详细信息：