diff --git a/src/sparseml/yolov8/trainers.py b/src/sparseml/yolov8/trainers.py index 5af0f97f22d..39a9346f381 100644 --- a/src/sparseml/yolov8/trainers.py +++ b/src/sparseml/yolov8/trainers.py @@ -29,6 +29,7 @@ from sparseml.optim.helpers import load_recipe_yaml_str from sparseml.pytorch.optim.manager import ScheduledModifierManager +from sparseml.pytorch.sparsification.quantization import skip_onnx_input_quantize from sparseml.pytorch.utils import ModuleExporter from sparseml.pytorch.utils.helpers import download_framework_model_by_recipe_type from sparseml.pytorch.utils.logger import LoggerManager, PythonLogger, WANDBLogger @@ -729,7 +730,13 @@ def export(self, **kwargs): else ["output0"], ) - onnx.checker.check_model(os.path.join(save_dir, name)) + complete_path = os.path.join(save_dir, name) + try: + skip_onnx_input_quantize(complete_path, complete_path) + except Exception: + pass + + onnx.checker.check_model(complete_path) deployment_folder = exporter.create_deployment_folder(onnx_model_name=name) if args["export_samples"]: trainer_config = get_cfg(cfg=DEFAULT_SPARSEML_CONFIG_PATH) diff --git a/src/sparseml/yolov8/utils/export_samples.py b/src/sparseml/yolov8/utils/export_samples.py index 9815fadf3ff..fb7a3094c1b 100644 --- a/src/sparseml/yolov8/utils/export_samples.py +++ b/src/sparseml/yolov8/utils/export_samples.py @@ -106,10 +106,15 @@ def export_sample_inputs_outputs( preprocessed_batch = preprocess(batch=batch, device=device) image = preprocessed_batch["img"] - # Save inputs as numpy array - _export_inputs(image, sample_in_dir, file_idx, save_inputs_as_uint8) # Save torch outputs as numpy array _export_torch_outputs(image, model, sample_out_dir_torch, file_idx) + + # Convert input data type if needed + if save_inputs_as_uint8: + image = (255 * image).to(dtype=torch.uint8) + + # Save inputs as numpy array + _export_inputs(image, sample_in_dir, file_idx) # Save onnxruntime outputs as numpy array _export_ort_outputs( image.cpu().numpy(), ort_session, sample_out_dir_ort, file_idx @@ -166,13 +171,9 @@ def _export_ort_outputs( numpy.savez(sample_output_filename, preds) -def _export_inputs( - image: torch.Tensor, sample_in_dir: str, file_idx: str, save_inputs_as_uint8: bool -): +def _export_inputs(image: torch.Tensor, sample_in_dir: str, file_idx: str): sample_in = image.detach().to("cpu") - if save_inputs_as_uint8: - sample_in = (255 * sample_in).to(dtype=torch.uint8) sample_input_filename = os.path.join(sample_in_dir, f"inp-{file_idx}.npz") numpy.savez(sample_input_filename, sample_in)