TorchView 使用教程

1. 项目介绍

TorchView 是一个用于可视化 PyTorch 模型的开源工具。它能够以图形的形式展示 PyTorch 模型的结构，包括张量、模块、torch.functions 等信息。TorchView 类似于 Keras 中的 plot_model 功能，但提供了更多的自定义选项和功能。

TorchView 的主要特点包括：

• 支持 PyTorch 版本 ≥ 1.7
• 提供多种可视化选项，如显示张量形状、隐藏内部张量等
• 支持元张量（meta tensors），避免在可视化过程中消耗大量内存
• 支持递归模块的展开和折叠

2. 项目快速启动

安装

首先，你需要安装 graphviz：

pip install graphviz

然后，安装 TorchView：

pip install torchview

或者，如果你想安装最新版本，可以直接从 GitHub 安装：

pip install git+https://github.com/mert-kurttutan/torchview.git

使用示例

以下是一个简单的使用示例，展示如何使用 TorchView 可视化一个简单的 PyTorch 模型：

from torchview import draw_graph
import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = MLP()
batch_size = 2

# 使用 TorchView 可视化模型
model_graph = draw_graph(model, input_size=(batch_size, 128), device='meta')
model_graph.visual_graph

3. 应用案例和最佳实践

案例1：可视化复杂的 ResNet 模型

以下代码展示了如何使用 TorchView 可视化一个复杂的 ResNet 模型，并展开嵌套模块：

import torchvision
from torchview import draw_graph

model = torchvision.models.resnet18()
model_graph = draw_graph(model, input_size=(1, 3, 224, 224), expand_nested=True)
model_graph.visual_graph

案例2：隐藏内部张量和函数

在某些情况下，你可能希望隐藏模型的内部张量和函数，以便更清晰地查看模型的结构。以下代码展示了如何实现这一点：

model_graph = draw_graph(MLP(), input_size=(2, 128), hide_inner_tensors=True, hide_module_functions=True)
model_graph.visual_graph

4. 典型生态项目

TensorBoard

TensorBoard 是另一个常用的 PyTorch 模型可视化工具，它提供了丰富的功能，包括模型结构可视化、训练过程中的指标监控等。你可以通过以下方式使用 TensorBoard 可视化模型：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("runs/mlp")
writer.add_graph(MLP(), torch.rand(2, 128))
writer.close()

Netron

Netron 是一个用于可视化深度学习模型的工具，支持多种模型格式，包括 ONNX、TensorFlow 等。你可以通过以下方式使用 Netron 可视化模型：

pip install netron
netron path/to/your/model.onnx

通过结合使用 TorchView 和其他生态项目，你可以更全面地理解和分析你的 PyTorch 模型。