forked from Purer-lyk/VisionWorkSpace
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pt_2onnx.py
58 lines (46 loc) · 1.72 KB
/
pt_2onnx.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import torch.onnx
from onnxmltools.utils.float16_converter import convert_float_to_float16
# from onnxconverter_common import convert_tensor_float_to_float16
import onnx, os
from v4.model_v4 import MultiTaskModel
fp16 = True
dynamic = True
file_name='2023_3_26_hj_num_1'
model_path = './weight/'+file_name+'.pt'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device('cpu')
def convert_fp16(save_path):
onnx_model = onnx.load_model(save_path)
fp16_onnx_model = convert_float_to_float16(onnx_model, keep_io_types=False)
onnx.save_model(fp16_onnx_model,save_path)
def export(model, save_path):
input_data = torch.randn(1, 1, 24, 24, device=device)
# 转化为onnx模型
input_names = ['input']
output_names = ['output']
torch.onnx.export(model, input_data, save_path, opset_version=9, verbose=True,
input_names=input_names, output_names=output_names,
dynamic_axes= {'input':{0:'1'}} if 'dyn' in save_path else {})
if 'fp16' in save_path:
convert_fp16(save_path)
def pt_2onnx():
print(device)
model = torch.load(model_path)
model.to(device)
model.eval()
initial_path = './onnx_model/' + file_name + '/'
if not os.path.exists(initial_path):
os.mkdir(initial_path)
initial_path += file_name
export(model, initial_path + '.onnx')
if dynamic:
save_path = initial_path + '_dyn'
export(model, save_path + '.onnx')
if fp16:
save_path = initial_path + '_fp16'
export(model, save_path + '.onnx')
if dynamic and fp16:
save_path = initial_path + '_dyn_fp16'
export(model, save_path + '.onnx')
if __name__ == "__main__":
pt_2onnx()