1 import os 2 import torch 3 import numpy as np 4 from Unet import UNET 5 os.environ["CUDA_VISIBLE_DEVICE"] = "" 6 7 def main(): 8 demo = Demo(model_path="/xxx.pth.tar", output="pathto/xxx.onnx") 9 demo.inference() 10 check_onnx(onnx_pth="path to xxx.onnx") 11 12 13 14 #检查onnx模型 15 def check_onnx(onnx_pth): 16 import onnx 17 #load the ONNX model 18 model = onnx.load(onnx_pth) 19 #check the IR is well formed 20 onnx.checker.check_model(model) 21 #print a human readable representation of graph 22 print(onnx.helper.printable_graph(model.graph)) 23 24 class WrappedModel(torch.nn.Module): 25 def __init__(self,model): 26 super().__init__() 27 self.model =model 28 29 def forward(self,x):31 outs=self.model(x) 32 new_outs=torch.sigmoid(outs) 33 return new_outs 34 35 36 class Demo(): 37 def __init__(self,model_path,output): 38 self.model_path =model_path 39 self.output_path = output 40 41 def init_torch_tensor(self): 42 self.device = 'cpu'#torch.device('cpu') 43 torch.set_default_tensor_type('torch.FloatTensor') 44 #use gpu or not 45 # if torch.cuda.is_available(): 46 # self.device = torch.device('cuda') 47 # torch.set_default_tensor_type('torch.FloatTensor') 48 # else: 49 # self.device = torch.device('cpu') 50 # torch.set_default_tensor_type('torch.FloatTensor') 51 52 def init_model(self,in_channels,out_channels): 53 model = UNET(in_channels=in_channels, out_channels=out_channels).to(self.device)#to('cuda') 54 return model 55 56 def resume(self, model, path): 57 if not os.path.exists(path): 58 print("Checkpoint not found:" + path) 59 return 60 states = torch.load(path, map_location=self.device)# 61 model.load_state_dict(states["state_dict"],strict=False)#states有两个key_value"state_dict","optimizer" 62 63 model_sig = WrappedModel(model) 64 print("Resume from " + path) 65 return model_sig 66 67 def inference(self): 68 #use gpu or cpu 69 self.init_torch_tensor() 70 #加载网络模型 71 model = self.init_model(in_channels=3,out_channels=2) 72 model_sig=self.resume(model, self.model_path) 73 #设置model的模式 74 model_sig.eval() 75 #设置输入 76 img = np.random.randint(0,255, size=(512,512,3),dtype=np.uint8) 77 img = img.astype(np.float32) 78 img = img / 255#(img / 255. - 0.5)/0.5 79 img = img.transpose((2,0,1)) #C H W 80 img = torch.from_numpy(img).unsqueeze(0).float() 81 #img = torch.randn(1,3,512,512) 82 ''' 83 设置动态可变维度 84 KEY(str) - 必须是input_names或output_names指定的名称，用来指定哪个变量需要使用到动态尺寸。 85 VALUE(dict or list) - 如果是一个dict，dict中的key是变量的某个维度，dict中的value是我们给这个维度取的名称。如果是一个list，则list中的元素都表示此变量的某个维度。 86 ''' 87 dynamic_axes = {'input':{0: 'batch_size', 2: 'height', 3: 'width'}, 88 'output': {0:'batch_size', 2: 'height', 3: 'width'}} 89 with torch.no_grad(): 90 img = img.to(self.device) 91 torch.onnx.export(model_sig, img, self.output_path, input_names=['input'], 92 output_names=['output'], dynamic_axes=dynamic_axes, keep_initializers_as_inputs=False,export_params=True, 93 verbose=True, opset_version=11) 94 95 if __name__ == '__main__': 96 main()