toolchaintechnologycenter/ModelTrainingPython
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects 1 Releases Wiki Activity

修复预测脚本和训练脚本的执行bug

Browse Source
This commit is contained in:
liyong
2025-10-21 19:49:21 +08:00
parent b9cce1d733
commit 4fb2da1366
10 changed files with 87 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
FC_ML_Data/FC_ML_Data_Output/Data_Output_Pytorch.py
View File

@@ -3,7 +3,7 @@ import torch
def export_model_pt(model,target,name = "model"):
script_model = torch.jit.script(model) # 或 torch.jit.trace(model, input)
script_model.save(target + name + ".pt")
script_model.save(target + name + ".pth")
#2 通用格式导出
def export_model_onnx(model,input_tensor,target,name="model"):
torch.onnx.export(model, input_tensor, target+ name + ".onnx")
@@ -11,12 +11,12 @@ def export_model_onnx(model,input_tensor,target,name="model"):
def export_model_bin(model,target,name = "weights"):
torch.save(model.state_dict(), target + name + ".bin")
def export_model(model,target,file_name,name):
def export_model(model,target,file_name,name,input_tensor):
if name == 'bin':
return export_model_bin(model,target,file_name)
if name == 'onnx':
return export_model_onnx(model,target,file_name)
if name == 'pt':
return export_model_bin(model,target,file_name)
return export_model_onnx(model,input_tensor,target,file_name)
if name == 'pth':
return export_model_pt(model,target,file_name)
else:
raise ValueError(f"不支持的导出类型")