diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000000..fd6f22b787e --- /dev/null +++ b/.dockerignore @@ -0,0 +1,7 @@ +* +!src/otx +!pyproject.toml +!README.md +!LICENSE +!MANIFEST.in +!docker/download_pretrained_weights.py diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda new file mode 100644 index 00000000000..56cffbbcd32 --- /dev/null +++ b/docker/Dockerfile.cuda @@ -0,0 +1,37 @@ +FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime AS base + +ARG http_proxy +ARG https_proxy +ARG no_proxy +ARG NON_ROOT_HOME=/home/non-root + +RUN apt-get update && apt-get install -y --no-install-recommends \ + libsm6=2:1.2.3-1 \ + libxext6=2:1.3.4-0ubuntu1 \ + ffmpeg=7:4.2.7-0ubuntu0.1 \ + libfontconfig1=2.13.1-2ubuntu3 \ + libxrender1=1:0.9.10-1 \ + libgl1-mesa-glx=21.2.6-0ubuntu0.1~20.04.2 \ + && rm -rf /var/lib/apt/lists/* + +RUN useradd -l -u 10001 non-root \ + && mkdir -p ${NON_ROOT_HOME} + +WORKDIR ${NON_ROOT_HOME} +COPY . src_dir +RUN chown -R non-root:non-root ${NON_ROOT_HOME} + +USER non-root + +ENV PATH=${PATH}:${NON_ROOT_HOME}/.local/bin + +RUN pip install --no-cache-dir src_dir/ && \ + otx install --do-not-install-torch && \ + rm -rf src_dir/ + +FROM base AS cuda + + +FROM base AS cuda_pretrained_ready +COPY docker/download_pretrained_weights.py download_pretrained_weights.py +RUN python download_pretrained_weights.py diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 00000000000..52a6942f5ef --- /dev/null +++ b/docker/README.md @@ -0,0 +1,24 @@ +# How to build cuda and cuda-pretrained-ready Docker images + +1. By executing the following commands, it will build two Docker images: `otx:${OTX_VERSION}-cuda` and `otx:${OTX_VERSION}-cuda-pretrained-ready`. + + ```console + git clone https://github.com/openvinotoolkit/training_extensions.git + cd docker + ./build.sh + ``` + +2. After that, you can check whether the images are built correctly such as + + ```console + docker image ls | grep otx + ``` + + Example: + + ```console + otx 2.0.0-cuda-pretrained-ready 4f3b5f98f97c 3 minutes ago 14.5GB + otx 2.0.0-cuda 8d14caccb29a 8 minutes ago 10.4GB + ``` + +`otx:${OTX_VERSION}-cuda` is a minimal Docker image where OTX is installed with CUDA supports. On the other hand, `otx:${OTX_VERSION}-cuda-pretrained-ready` includes all the model pre-trained weights that OTX provides in addition to `otx:${OTX_VERSION}-cuda`. diff --git a/docker/build.sh b/docker/build.sh new file mode 100755 index 00000000000..5888516d481 --- /dev/null +++ b/docker/build.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# shellcheck disable=SC2154 + +OTX_VERSION=$(python -c 'import otx; print(otx.__version__)') +THIS_DIR=$(dirname "$0") + +echo "Build OTX ${OTX_VERSION} CUDA Docker image..." +docker build \ + --build-arg http_proxy="${http_proxy}" \ + --build-arg https_proxy="${https_proxy}" \ + --build-arg no_proxy="${no_proxy}" \ + --target cuda \ + -t "otx:${OTX_VERSION}-cuda" \ + -f "${THIS_DIR}/Dockerfile.cuda" "${THIS_DIR}"/.. + +echo "Build OTX ${OTX_VERSION} CUDA pretrained-ready Docker image..." +docker build \ + --build-arg http_proxy="${http_proxy}" \ + --build-arg https_proxy="${https_proxy}" \ + --build-arg no_proxy="${no_proxy}" \ + --target cuda_pretrained_ready \ + -t "otx:${OTX_VERSION}-cuda-pretrained-ready" \ + -f "${THIS_DIR}/Dockerfile.cuda" "${THIS_DIR}"/.. diff --git a/docker/download_pretrained_weights.py b/docker/download_pretrained_weights.py new file mode 100644 index 00000000000..7fee845239e --- /dev/null +++ b/docker/download_pretrained_weights.py @@ -0,0 +1,52 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Helper script to download all the model pre-trained weights.""" + +import logging +from pathlib import Path + +from importlib_resources import files +from omegaconf import OmegaConf +from otx.core.utils.instantiators import partial_instantiate_class + +logging.basicConfig( + level=logging.INFO, + filename="download_pretrained_weights.log", + filemode="w", +) + +logger = logging.getLogger() + + +def download_all() -> None: + """Download pre-trained weights of all models.""" + recipe_dir = Path(files("otx") / "recipe") + + for config_path in recipe_dir.glob("**/*.yaml"): + if "_base_" in str(config_path): + msg = f"Skip {config_path} since it is a base config." + logger.warning(msg) + continue + if config_path.name == "openvino_model.yaml": + msg = f"Skip {config_path} since it is not a PyTorch config." + logger.warning(msg) + continue + if "anomaly_" in str(config_path) or "otx_dino_v2" in str(config_path) or "h_label_cls" in str(config_path): + msg = f"Skip {config_path} since those models show errors on instantiation." + logger.warning(msg) + continue + + config = OmegaConf.load(config_path) + init_model = next(iter(partial_instantiate_class(config.model))) + try: + model = init_model() + msg = f"Downloaded pre-trained model weight of {model!s}" + logger.info(msg) + except Exception: + msg = f"Error on instiating {config_path}" + logger.exception(msg) + + +if __name__ == "__main__": + download_all() diff --git a/pyproject.toml b/pyproject.toml index d5c222cd838..cead5265c8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,6 +87,9 @@ mmlab = [ # This causes an error when training the mm model, so fix the version first. "oss2==2.17.0", ] +anomaly = [ + "anomalib==1.0.0", +] [project.scripts] otx = "otx.cli:main" diff --git a/src/otx/cli/install.py b/src/otx/cli/install.py index 607d98863f1..37523539b87 100644 --- a/src/otx/cli/install.py +++ b/src/otx/cli/install.py @@ -59,10 +59,15 @@ def add_install_parser(subcommands_action: _ActionSubCommands) -> None: help="Set Logger level to INFO", action="store_true", ) + parser.add_argument( + "--do-not-install-torch", + help="Do not install PyTorch. Choose this option if you already install PyTorch.", + action="store_true", + ) subcommands_action.add_subcommand("install", parser, help="Install OTX requirements.") -def otx_install(option: str | None = None, verbose: bool = False) -> int: +def otx_install(option: str | None = None, verbose: bool = False, do_not_install_torch: bool = False) -> int: """Install OTX requirements. Args: @@ -92,12 +97,15 @@ def otx_install(option: str | None = None, verbose: bool = False) -> int: # This is done to parse the correct version of torch (cpu/cuda) and mmcv (mmcv/mmcv-full). torch_requirement, mmcv_requirements, other_requirements = parse_requirements(requirements) - # Get install args for torch to install it from a specific index-url install_args: list[str] = [] - torch_install_args = get_torch_install_args(torch_requirement) # Combine torch and other requirements. - install_args = other_requirements + torch_install_args + install_args = ( + # Get install args for torch to install it from a specific index-url + other_requirements + get_torch_install_args(torch_requirement) + if not do_not_install_torch + else other_requirements + ) # Parse mmX requirements if the task requires mmX packages. mmcv_install_args = [] diff --git a/tests/unit/cli/test_install.py b/tests/unit/cli/test_install.py index 2f8d0750538..a55c587a8a9 100644 --- a/tests/unit/cli/test_install.py +++ b/tests/unit/cli/test_install.py @@ -28,7 +28,7 @@ def test_add_install_parser(self) -> None: assert parser_subcommands.choices.get("install") is not None install_parser = parser_subcommands.choices.get("install") argument_list = [action.dest for action in install_parser._actions] - expected_argument = ["help", "option", "verbose"] + expected_argument = ["help", "option", "verbose", "do_not_install_torch"] assert argument_list == expected_argument def test_install_extra(self, mocker: MockerFixture) -> None: