diff --git a/.github/workflows/test_gpu.yaml b/.github/workflows/test_gpu.yaml index 661aa164..d76f763f 100644 --- a/.github/workflows/test_gpu.yaml +++ b/.github/workflows/test_gpu.yaml @@ -27,8 +27,7 @@ jobs: uses: actions/checkout@v2 - name: Build GPU Docker image - if: github.event_name == 'push' && (contains(github.event_path, 'docker/gpu.dockerfile') || contains(github.event_path, 'gpu_requirements.txt')) - run: docker build -f docker/gpu.dockerfile -t optimum-benchmark-gpu . + run: docker build --no-cache --build-arg CACHEBUST=$(date +%s) -f docker/gpu.dockerfile -t optimum-benchmark-gpu . - name: Run GPU tests run: docker run diff --git a/.github/workflows/test_ort_training.yaml b/.github/workflows/test_ort_training.yaml index 4166d3a9..6540628e 100644 --- a/.github/workflows/test_ort_training.yaml +++ b/.github/workflows/test_ort_training.yaml @@ -27,8 +27,7 @@ jobs: uses: actions/checkout@v2 - name: Build OnnxRuntime Training Docker image - if: github.event_name == 'push' && (contains(github.event_path, 'docker/ort_training.dockerfile') || contains(github.event_path, 'gpu_requirements.txt')) - run: docker build -f docker/ort_training.dockerfile -t optimum-benchmark-ort-training . + run: docker build --no-cache --build-arg CACHEBUST=$(date +%s) -f docker/ort_training.dockerfile -t optimum-benchmark-ort-training . - name: Run OnnxRuntime Training tests run: docker run diff --git a/docker/gpu.dockerfile b/docker/gpu.dockerfile index 332c6f0e..f01d2bbe 100644 --- a/docker/gpu.dockerfile +++ b/docker/gpu.dockerfile @@ -29,6 +29,9 @@ RUN unattended-upgrade RUN apt-get autoremove -y RUN pip install --upgrade pip +# this line forces the docker build to rebuild from this point on +ARG CACHEBUST=1 + # Install optimum-benchmark dependencies COPY gpu_requirements.txt /tmp/gpu_requirements.txt RUN pip install -r /tmp/gpu_requirements.txt \ No newline at end of file diff --git a/docker/ort_training.dockerfile b/docker/ort_training.dockerfile index 6e07c3a7..c40654d8 100644 --- a/docker/ort_training.dockerfile +++ b/docker/ort_training.dockerfile @@ -64,6 +64,9 @@ ENV TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" RUN $PYTHON_EXE -m pip install --upgrade protobuf==3.20.2 RUN $PYTHON_EXE -m torch_ort.configure +# this line forces the docker build to rebuild from this point on +ARG CACHEBUST=1 + # Install optimum-benchmark dependencies COPY requirements.txt /tmp/requirements.txt RUN pip install -r /tmp/requirements.txt \ No newline at end of file diff --git a/docker/scripts/build_gpu.sh b/docker/scripts/build_gpu.sh index b1591e8f..5ecfdb6a 100644 --- a/docker/scripts/build_gpu.sh +++ b/docker/scripts/build_gpu.sh @@ -1 +1 @@ -docker build -f docker/gpu.dockerfile -t optimum-benchmark-gpu . \ No newline at end of file +docker build --no-cache --build-arg CACHEBUST=$(date +%s) -f docker/gpu.dockerfile -t optimum-benchmark-gpu . \ No newline at end of file diff --git a/docker/scripts/build_ort_training.sh b/docker/scripts/build_ort_training.sh index 24011ce7..d2efd477 100644 --- a/docker/scripts/build_ort_training.sh +++ b/docker/scripts/build_ort_training.sh @@ -1 +1 @@ -docker build -f docker/ort_training.dockerfile -t optimum-benchmark-ort-training . \ No newline at end of file +docker build --no-cache --build-arg CACHEBUST=$(date +%s) -f docker/ort_training.dockerfile -t optimum-benchmark-ort-training . \ No newline at end of file diff --git a/optimum_benchmark/backends/optimum_utils.py b/optimum_benchmark/backends/optimum_utils.py index 71f58a9f..dbbbe84b 100644 --- a/optimum_benchmark/backends/optimum_utils.py +++ b/optimum_benchmark/backends/optimum_utils.py @@ -156,8 +156,11 @@ def main_export( if fp16 is True and device == "cpu": raise ValueError( - "The --fp16 option is supported only when exporting on GPU. Please pass the option `--device cuda`." + "FP16 export is supported only when exporting on GPU. Please pass the option `--device cuda`." ) + float_dtype = "fp16" + else: + float_dtype = "fp32" output = Path(output) if not output.exists(): @@ -277,6 +280,7 @@ def main_export( monolith=monolith, custom_onnx_configs=custom_onnx_configs if custom_onnx_configs is not None else {}, custom_architecture=custom_architecture, + float_dtype=float_dtype, fn_get_submodels=fn_get_submodels, preprocessors=preprocessors, _variant=_variant, @@ -328,7 +332,7 @@ def main_export( f" referring to `optimum.exporters.tasks.TaskManager`'s `_TASKS_TO_AUTOMODELS`." ) - onnx_files_subpaths = None + onnx_files_subpaths = [key + ".onnx" for key in models_and_onnx_configs.keys()] else: # save the subcomponent configuration for model_name in models_and_onnx_configs: @@ -373,67 +377,3 @@ def main_export( models_and_onnx_configs[key] = ("dummy_model", models_and_onnx_configs[key][1]) return onnx_config, models_and_onnx_configs - - # if optimize is not None: - # from ...onnxruntime import AutoOptimizationConfig, ORTOptimizer - - # if onnx_files_subpaths is None: - # onnx_files_subpaths = [key + ".onnx" for key in models_and_onnx_configs.keys()] - # optimizer = ORTOptimizer.from_pretrained(output, file_names=onnx_files_subpaths) - - # optimization_config = AutoOptimizationConfig.with_optimization_level(optimization_level=optimize) - - # optimization_config.disable_shape_inference = True - # optimizer.optimize(save_dir=output, optimization_config=optimization_config, file_suffix="") - - # # Optionally post process the obtained ONNX file(s), for example to merge the decoder / decoder with past if any - # # TODO: treating stable diffusion separately is quite ugly - # if not no_post_process and not is_stable_diffusion: - # try: - # logger.info("Post-processing the exported models...") - # models_and_onnx_configs, onnx_files_subpaths = onnx_config.post_process_exported_models( - # output, models_and_onnx_configs, onnx_files_subpaths - # ) - # except Exception as e: - # raise Exception( - # f"The post-processing of the ONNX export failed. The export can still be performed by passing the option --no-post-process. Detailed error: {e}" - # ) - - # if is_stable_diffusion: - # use_subprocess = ( - # False # TODO: fix Can't pickle local object 'get_stable_diffusion_models_for_export..' - # ) - # elif model.config.model_type in UNPICKABLE_ARCHS: - # # Pickling is bugged for nn.utils.weight_norm: https://github.com/pytorch/pytorch/issues/102983 - # # TODO: fix "Cowardly refusing to serialize non-leaf tensor" error for wav2vec2-conformer - # use_subprocess = False - - # if do_validation is True: - # try: - # validate_models_outputs( - # models_and_onnx_configs=models_and_onnx_configs, - # onnx_named_outputs=onnx_outputs, - # atol=atol, - # output_dir=output, - # onnx_files_subpaths=onnx_files_subpaths, - # input_shapes=input_shapes, - # device=device, - # dtype=torch_dtype, - # use_subprocess=use_subprocess, - # model_kwargs=model_kwargs, - # ) - # logger.info(f"The ONNX export succeeded and the exported model was saved at: {output.as_posix()}") - # except ShapeError as e: - # raise e - # except AtolError as e: - # logger.warning( - # f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {output.as_posix()}" - # ) - # except OutputMatchError as e: - # logger.warning( - # f"The ONNX export succeeded with the warning: {e}.\n The exported model was saved at: {output.as_posix()}" - # ) - # except Exception as e: - # raise Exception( - # f"An error occured during validation, but the model was saved nonetheless at {output.as_posix()}. Detailed error: {e}." - # )