部分节点访问与数据比对

master
wangchongwu 3 years ago
parent 55b05c2836
commit b10e13caea

@ -5,6 +5,13 @@ import matplotlib.pyplot as plt
import torch import torch
import cv2 import cv2
import torchvision.transforms as transforms import torchvision.transforms as transforms
import numpy as np
import tkinter
import tkinter.messagebox #弹窗库
from PyQt5.QtWidgets import *
#定义网络 #定义网络
class Net(nn.Module): class Net(nn.Module):
def __init__(self): def __init__(self):
@ -35,6 +42,10 @@ net.to(device)
net.eval() net.eval()
for name in net.state_dict():
print(name)
print(net.state_dict()['conv1.weight'])
img = cv2.imread('./data/8.png') img = cv2.imread('./data/8.png')
img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY) img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
img = img/255 img = img/255
@ -49,9 +60,35 @@ out = net(imgTensor)
print(out.data.max(1, keepdim=True)[1][0].item()) print(out.data.max(1, keepdim=True)[1][0].item())
tkinter.messagebox.showinfo('pt推理结果', out.data.max(1, keepdim=True)[1][0].item())
# onnx model
dummy_input = torch.randn(1, 1,28,28).to(device)#输入大小 #data type nchw dummy_input = torch.randn(1, 1,28,28).to(device)#输入大小 #data type nchw
torch.onnx.export(net, dummy_input, "ministNet.onnx", verbose=True, input_names=['input_111'], output_names=['output_111']) torch.onnx.export(net, dummy_input, "ministNet.onnx", verbose=True, input_names=['input_111'], output_names=['output_111'])
net = cv2.dnn.readNetFromONNX("ministNet.onnx") # 加载训练好的识别模型
img = cv2.imread('./data/8.png')
img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
img = img/255
transf = transforms.ToTensor()
imgTensor = transf(img).unsqueeze(0)
im = img[np.newaxis, np.newaxis,:, :]
im = im.astype(np.float32)
outNames = net.getUnconnectedOutLayersNames()
net.setInput(im)
out = net.forward(outNames)
ind = np.where(out[0][0]==np.max(out[0][0]))
print(ind[0])
tkinter.messagebox.showinfo('onnx推理结果',ind[0])

@ -1,27 +0,0 @@
import cv2
import torchvision.transforms as transforms
import numpy as np
net = cv2.dnn.readNetFromONNX("ministNet.onnx") # 加载训练好的识别模型
img = cv2.imread('./data/8.png')
img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
img = img/255
transf = transforms.ToTensor()
imgTensor = transf(img).unsqueeze(0)
im = img[np.newaxis, np.newaxis,:, :]
im = im.astype(np.float32)
outNames = net.getUnconnectedOutLayersNames()
net.setInput(im)
out = net.forward(outNames)
ind = np.where(out[0][0]==np.max(out[0][0]))
print(ind[0])

Binary file not shown.

Binary file not shown.

Binary file not shown.
Loading…
Cancel
Save