神经网络可视化工具Zetane的使用方法

最新推荐文章于 2024-09-19 07:01:42 发布

uestc_Venn

最新推荐文章于 2024-09-19 07:01:42 发布

阅读量610

点赞数

分类专栏：人工智能神经网络可视化机器学习文章标签：神经网络深度学习人工智能

本文链接：https://blog.csdn.net/Mathematic_Van/article/details/132638992

版权

人工智能同时被 3 个专栏收录

14 篇文章 0 订阅

订阅专栏

机器学习

4 篇文章 0 订阅

订阅专栏

神经网络可视化

3 篇文章 0 订阅

订阅专栏

本文介绍了如何将PyTorch模型（如ResNet18）的.pth文件转换为ONNX格式，包括设置模型、数据预处理、使用torch.onnx.export进行导出，以及安装ONNX模块并提供上传模型的示例。

摘要由CSDN通过智能技术生成

安装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)