模型转换
pytorch转onnx
import torch.utils.data
from torch.autograd import Variable
from squeezenet import squeezenet1_2
import onnxruntime
import numpy as np
from onnxruntime.datasets import get_example
import cv2
import onnx
model_file = '/home/bbt/qinghua/detetion/pytorch-mobilenet/model/model_best_squeezenet1_2_age0.926.pth.tar'
num_class =3
# create model
model=squeezenet1_2(pretrained=False,num_classes=num_class)
model = torch.nn.DataParallel(model)
# optionally resume from a checkpoint
checkpoint = torch.load(model_file)
model.load_state_dict(checkpoint['state_dict'])
model.to('cpu')
model.eval()
# model.cpu()
#accuracy(mode)
# dummy_input = Variable(torch.randn(1, 3, 224, 224))
input=cv2.imread('/home/bbt/age00.jpg')
input=cv2.resize(input,(224,224))
input=np.transpose(input, (2, 0, 1)).astype(np.float32)
now_image1 = Variable(torch.