安装Zetane

模型文件格式转换——.pth文件转换为ONNX文件

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torchvision.models import resnet18

# 设置模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = resnet18(pretrained=False, num_classes=10)
model.conv1 = nn.Conv2d(3, 64, 3, stride=1, padding=1, bias=False)
model.maxpool = nn.MaxPool2d(1, 1, 0)
model.load_state_dict(torch.load('D:\可视化\model.pth'))
model = model.to(device)
model.eval()


# 数据转换
transform = transforms.Compose([
    transforms.Resize(32),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True)

# 获取一张图片
image, _ = next(iter(train_loader))

# 导出模型
onnx_path = "D:\可视化\model.onnx"
torch.onnx.export(model, image.to(device), onnx_path, verbose=True)

安装ONNX模块清华源镜像

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnx

上传模型文件（支持ONNX，Keras (.h5)和ZTN文件格式）

使用示例

• 右键拖动画面或视野维度
• 鼠标滚轮缩放画面