pth和onnx推理结果对比代码
需积分: 5 112 浏览量
2023-02-16
15:01:18
上传
评论 1
收藏 36KB DOCX 举报
import torch
import torchvision.models as models
import torchvision.transforms as transforms
import onnxruntime as ort
from PIL import Image
def pth_to_onnx():
model = models.resnet18()
model.load_state_dict(torch.load("F:/work/Atlantic/resnet18.pth",
map_location='cuda'))
model.eval()
input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, input, "resnet18.onnx", verbose=True,
opset_version=11, export_params=True)
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad
else tensor.cpu().numpy()
def get_test_transform():
return transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor()
])
if __name__ == '__main__':
pth_to_onnx()
#y 验证
img_input = Image.open("img.png")
img = get_test_transform()(img_input)
img = img.unsqueeze_(0)
# dummy_input = torch.randn(10, 3, 224, 224)
pth_model_path = "F:/work/Atlantic/resnet18.pth"
## pth 测试
model_net = models.resnet18()
model_net.load_state_dict(torch.load(pth_model_path))
model_net.eval()
output = model_net(img)
print("pth weights", output.detach().cpu().numpy())
print("pth prediction", output.argmax(dim=1)[0].item())#pth
prediction 640
###onnx 测试
ort_session = ort.InferenceSession('resnet18.onnx')
inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}
outputs = ort_session.run(None, inputs)[0]
print("onnx weights", outputs)
剩余12页未读，继续阅读
资源评论
