pytorch 命令行运行_pytorch模型可视化

上班时间摸鱼写个博客。

搞了4.5个小时的Pytorch模型可视化：

一.netron

1.命令行端口模式

d = torch.rand(1, 3, 416, 416)

m = model()

o = m(d)

onnx_path = "onnx_model_name.onnx"

torch.onnx.export(m, d, onnx_path)

netron.start(onnx_path)

2.软件客户端模式

import torch.nn as nn

import torch.nn.functional as F

import torch

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

class Net(nn.Module):

def __init__(self):

super(Net, self).__init__()

self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)

self.conv2 = nn.Conv2d(10, 20, 5)

self.conv3 = nn.Conv2d(20, 40, 3)

self.mp = nn.MaxPool2d(2)

self.mp1 = nn.MaxPool2d(2)

self.fc1 = nn.Linear(2560, 512)

self.fc2 = nn.Line