一、TessorBoard
二、netron查看模型结构

pip install torch torchvision tensorboard
#打开一个终端执行，然后访问http://localhost:6006
tensorboard --logdir runs

TensorBoard记录：

Loss曲线
Accuracy曲线
输入图片
网络结构 （复杂的一般不用这个）
网络结构
权重变化

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

from torch.utils.tensorboard import SummaryWriter


# ======================
# 1. TensorBoard
# ======================

writer = SummaryWriter("runs/cnn_demo")


# ======================
# 2. 数据集 MNIST
# ======================

transform = transforms.ToTensor()


train_dataset = torchvision.datasets.MNIST(
    root="./data",
    train=True,
    download=True,
    transform=transform
)


test_dataset = torchvision.datasets.MNIST(
    root="./data",
    train=False,
    download=True,
    transform=transform
)


train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True
)


test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False
)


# ======================
# 3. CNN模型
# ======================

class CNN(nn.Module):

    def __init__(self):
        super().__init__()

        self.conv = nn.Sequential(

            nn.Conv2d(
                1,      # 输入通道
                8,      # 输出通道
                3,
                padding=1
            ),

            nn.ReLU(),

            nn.MaxPool2d(2),


            nn.Conv2d(
                8,
                16,
                3,
                padding=1
            ),

            nn.ReLU(),

            nn.MaxPool2d(2)
        )


        self.fc = nn.Sequential(

            nn.Flatten(),

            nn.Linear(
                16*7*7,
                10
            )
        )


    def forward(self,x):

        x=self.conv(x)

        x=self.fc(x)

        return x



# ======================
# 4. 创建模型
# ======================

device = "cuda" if torch.cuda.is_available() else "cpu"

model=CNN().to(device)


loss_fn=nn.CrossEntropyLoss()


optimizer=torch.optim.Adam(
    model.parameters(),
    lr=0.001
)



# ======================
# 5. TensorBoard记录网络结构
# ======================

dummy=torch.randn(
    1,
    1,
    28,
    28
).to(device)


writer.add_graph(
    model,
    dummy
)



# ======================
# 6. 训练
# ======================

step=0


epochs=1


for epoch in range(epochs):

    model.train()

    total_loss=0


    for images,labels in train_loader:


        images=images.to(device)
        labels=labels.to(device)


        # forward
        outputs=model(images)


        loss=loss_fn(
            outputs,
            labels
        )


        # backward

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()



        # ===== TensorBoard Loss =====

        writer.add_scalar(
            "train/loss",
            loss.item(),
            step
        )


        step+=1

        total_loss += loss.item()



    # ======================
    # 测试准确率
    # ======================

    model.eval()


    correct=0
    total=0


    with torch.no_grad():

        for images,labels in test_loader:

            images=images.to(device)
            labels=labels.to(device)


            outputs=model(images)


            pred=torch.argmax(
                outputs,
                dim=1
            )


            correct += (
                pred==labels
            ).sum().item()


            total += labels.size(0)



    acc=correct/total


    print(
        f"Epoch {epoch+1}, "
        f"loss={total_loss:.3f}, "
        f"acc={acc:.4f}"
    )



    # ===== TensorBoard Accuracy =====

    writer.add_scalar(
        "test/accuracy",
        acc,
        epoch
    )


    # ===== TensorBoard 权重 =====

    writer.add_histogram(
        "conv1_weight",
        model.conv[0].weight,
        epoch
    )



    # ===== TensorBoard 图片 =====

    images,_=next(iter(train_loader))


    writer.add_images(
        "input_images",
        images[:16],
        epoch
    )



# ======================
# 7. 保存模型
# ======================

torch.save(
    model.state_dict(),
    "cnn_mnist.pth"
)


writer.close()


print("训练完成")

2.netron模型结构查看

pip install netron onnx
PyTorch模型
    ↓
导出ONNX
    ↓
用Netron打开
netron cnn_model.onnx
然后浏览器会打开模型结构页面
http://localhost:8080

以下是普通模型导出为onnx的代码：

import torch

model.eval()

example_input = torch.randn(
    1, 1, 28, 28
).to(device)

torch.onnx.export(
    model,
    example_input,
    "cnn_model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=18,
    dynamo=True
)

print("已导出 cnn_model.onnx")