onnx测试简化

╰+哭是因爲堅強的太久メ 2022-09-09 06:45 199阅读 0赞

if img_path is None:
    		        img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.uint8)
    		    else:
    		        img = cv2.imread(img_path) #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    		        img = cv2.resize(img,(112,112))
    		        
    		        # img -=(0.485*255, 0.456*255, 0.406*255)
    		        # img /= (0.229*255, 0.224*255, 0.225*255)
    		        
    		        img -=(0.5*255, 0.5*255, 0.5*255)
    		        img /= (0.5*255, 0.5*255, 0.5*255)
    		        img = img.transpose(2, 0, 1)  #hwc to chw
    		        image = np.array(img)[np.newaxis, :, :, :].astype(np.float32)  # chw  to  nchw
    
                import onnxruntime as rt
                sess = rt.InferenceSession("./res50cbam_sim.onnx")
                input_name1 = [input.name for input in sess.get_inputs()][0]
                outname = [output.name for output in sess.get_outputs()]
                output = sess.run(outname, {input_name1: image.numpy()})
                print(output[0].shape)
                print(output[0][0][0:10])
                print("=================>")

if img_path is None:
    		        img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.uint8)
    		    else:
    		        img = cv2.imread(img_path) #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    		        img = cv2.resize(img,(112,112))
    		        
    		        # img -=(0.485*255, 0.456*255, 0.406*255)
    		        # img /= (0.229*255, 0.224*255, 0.225*255)
    		        
    		        img -=(0.5*255, 0.5*255, 0.5*255)
    		        img /= (0.5*255, 0.5*255, 0.5*255)
    		        img = img.transpose(2, 0, 1)  #hwc to chw
    		        image = np.array(img)[np.newaxis, :, :, :].astype(np.float32)  # chw  to  nchw
    
                print("=================>")
                print(image.shape)
                torch.onnx.export(self.model, image, "./res50cbam.onnx",verbose=True,opset_version=11) 
                from onnxsim import simplify  # pip install onnx-simplifier
                import onnx
                model = onnx.load("./res50cbam.onnx")
                onnx.checker.check_model(model)
                model, check = simplify(model)
                assert check, "Simplified ONNX model could not be validated"
                onnx.save(model, "./res50cbam_sim.onnx")