diff --git a/routers/repo/ai_model_convert.go b/routers/repo/ai_model_convert.go index aac66d078..ba98e5b7a 100644 --- a/routers/repo/ai_model_convert.go +++ b/routers/repo/ai_model_convert.go @@ -52,6 +52,9 @@ const ( CONVERT_FORMAT_ONNX = 0 CONVERT_FORMAT_TRT = 1 + NetOutputFormat_FP32 = 0 + NetOutputFormat_FP16 = 1 + NPU_MINDSPORE_IMAGE_ID = 35 NPU_TENSORFLOW_IMAGE_ID = 121 @@ -460,7 +463,15 @@ func getGpuModelConvertCommand(name string, modelFile string, modelConvert *mode h = intputshape[2] w = intputshape[3] } - command += "python3 /code/" + bootfile + " --model " + modelFile + " --n " + n + " --c " + c + " --h " + h + " --w " + w + " > " + ModelMountPath + "/" + name + "-" + LogFile + command += "python3 /code/" + bootfile + " --model " + modelFile + " --n " + n + " --c " + c + " --h " + h + " --w " + w + if modelConvert.DestFormat == CONVERT_FORMAT_TRT { + if modelConvert.NetOutputFormat == NetOutputFormat_FP16 { + command += " --fp16 True" + } else { + command += " --fp16 False" + } + } + command += " > " + ModelMountPath + "/" + name + "-" + LogFile return command }