Open source detectron2 handler and trainer scripts and dockerfiles. (#2286)

2026-09-26 14:42:04 +00:00 · 2023-09-28 15:23:24 +00:00
parent eb93058f8e
commit 572cbaf8b7
5 changed files with 441 additions and 0 deletions
@@ -0,0 +1,90 @@
+# Dockerfile for Detectron2 serving.
+#
+# To build:
+# docker build -f model_oss/detectron2/dockerfile/serving.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+FROM pytorch/torchserve:0.7.0-cpu
+
+USER root
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# run and update some basic packages software packages, including security libs
+RUN apt-get update &&  apt-get install -y \
+    software-properties-common && \
+    add-apt-repository -y ppa:ubuntu-toolchain-r/test && \
+    apt-get update && apt-get install -y \
+    gcc-9 g++-9 apt-transport-https ca-certificates gnupg curl
+
+# Install gcloud tools for gsutil as well as debugging
+RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | \
+    tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | \
+    apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
+    apt-get update -y && apt-get install google-cloud-sdk -y
+
+USER model-server
+
+# install detectron2 dependencies
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install --user numpy==1.24.2
+RUN python3 -m pip install --user opencv-python==4.7.0.72
+RUN python3 -m pip install --user 'git+https://github.com/facebookresearch/detectron2.git@v0.6'
+
+# Install GCS storage library.
+RUN pip install google-cloud-storage==2.6.0
+
+# For mask encoding.
+RUN pip install --upgrade pycocotools==2.0.6
+
+ARG MODEL_NAME=detectron2_serving
+ENV MODEL_NAME="${MODEL_NAME}"
+
+# health and prediction listener ports
+ARG AIP_HTTP_PORT=7080
+ENV AIP_HTTP_PORT="${AIP_HTTP_PORT}"
+
+ARG MODEL_MGMT_PORT=7081
+
+# expose health and prediction listener ports from the image
+EXPOSE "${AIP_HTTP_PORT}"
+EXPOSE "${MODEL_MGMT_PORT}"
+EXPOSE 8080 8081 8082 7070 7071
+
+# create torchserve configuration file
+USER root
+RUN echo "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${AIP_HTTP_PORT}\n"  \
+    "management_address=http://0.0.0.0:${MODEL_MGMT_PORT}" >> /home/model-server/config.properties
+USER model-server
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY ./model_oss/detectron2/handler.py /home/model-server/handler.py
+WORKDIR /home/model-server/
+
+# Create model archive file packaging model artifacts and dependencies.
+# Note(lavrai): The model `.pth` file and `cfg.yaml` file will be set by the
+# customer as an environment variable and will be later loaded by the
+# `handler.py` file.
+RUN torch-model-archiver \
+  --model-name="${MODEL_NAME}" \
+  --version=1.0 \
+  --handler=/home/model-server/handler.py \
+  --export-path=/home/model-server/model-store \
+  -f
+
+# run Torchserve HTTP serve to respond to prediction requests
+CMD ["ls", "-ltr", "/home/model-server/model-store/", ";", \
+    "torchserve", "--start", "--ts-config=/home/model-server/config.properties", \
+    "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", \
+    "--model-store", "/home/model-server/model-store"]
@@ -0,0 +1,100 @@
+# Dockerfile for Detectron2 training.
+#
+# To build:
+# docker build -f model_oss/detectron2/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM nvidia/cuda:11.1.1-cudnn8-devel-ubuntu18.04
+# Using an older system (18.04) to avoid opencv incompatibility (issue#3524).
+
+ENV DEBIAN_FRONTEND noninteractive
+RUN apt-get update && apt-get install -y \
+    python3.7 python3.7-dev python3.7-distutils \
+    python3-opencv ca-certificates git wget sudo ninja-build \
+    curl wget vim
+
+# Make python3 available for python3.7.
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.6 1
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.7 2
+RUN update-alternatives --config python3
+# Make python available for python3.7.
+RUN ln -sv /usr/bin/python3.7 /usr/bin/python
+
+# Create a non-root user.
+ARG USER_ID=1000
+RUN useradd -m --no-log-init --system  --uid ${USER_ID} appuser -g sudo
+RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
+USER appuser
+WORKDIR /home/appuser
+
+ENV PATH="/home/appuser/.local/bin:${PATH}"
+RUN wget https://bootstrap.pypa.io/pip/get-pip.py && \
+    python3.7 get-pip.py --user && \
+    rm get-pip.py
+
+# Important! Otherwise, it uses existing numpy from host-modules
+# which throws error.
+RUN pip install --user numpy==1.20.3
+
+# Install dependencies:
+# See https://pytorch.org/ for other options if you use
+# a different version of CUDA.
+RUN pip install --user tensorboard==2.11.0
+# cmake from apt-get is too old.
+RUN pip install --user cmake==3.25.2
+RUN pip install --user torch==1.10.0+cu111 torchvision==0.11.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
+RUN pip install --user setuptools==59.5.0
+RUN pip install --user opencv-python==4.7.0.72
+RUN pip install --user cloudml-hypertune==0.1.0.dev6
+RUN pip install --user fvcore==0.1.5.post20221221
+# Install detectron2.
+RUN git clone -b v0.6 https://github.com/facebookresearch/detectron2 detectron2_repo
+# Set FORCE_CUDA because during `docker build` cuda is not accessible.
+ENV FORCE_CUDA="1"
+# This will by default build detectron2 for all common cuda
+# architectures and take a lot more time,
+# because inside `docker build`, there is no way to tell
+# which architecture will be used.
+ARG TORCH_CUDA_ARCH_LIST="Kepler;Kepler+Tesla;Maxwell;Maxwell+Tegra;Pascal;Volta;Turing"
+ENV TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}"
+RUN pip install --user -e detectron2_repo
+
+# Set a fixed model cache directory.
+ENV FVCORE_CACHE="/tmp"
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model-garden detectron2 files to '/home/appuser/trainer' folder.
+ADD ./model_oss/detectron2 /home/appuser/trainer
+
+################ Copy plain_train_net.py to task.py and
+# then modify it using sed commands. ###################
+# Src: https://github.com/facebookresearch/detectron2/blob/v0.6/tools/plain_train_net.py
+RUN sudo cp /home/appuser/detectron2_repo/tools/plain_train_net.py /home/appuser/trainer/task.py
+# Make additional changes to task.py.
+# Note(lavrai): Start adding SED commands from end of file towards the top
+# so that the line numbers do not keep changing for the source file.
+# For entry-point:
+RUN sudo sed -i "214 d" /home/appuser/trainer/task.py
+RUN sudo sed -i "213 a\    default_arg_parser = default_argument_parser()" /home/appuser/trainer/task.py
+RUN sudo sed -i "214 a\    extended_parser = trainer_utils.extend_parser_arguments(default_arg_parser)" /home/appuser/trainer/task.py
+RUN sudo sed -i "215 a\    args = extended_parser.parse_args()" /home/appuser/trainer/task.py
+# For main() function:
+RUN sudo sed -i "192 a\    trainer_utils.register_dataset(args)" /home/appuser/trainer/task.py
+# For setup() function:
+RUN sudo sed -i "184 a\    cfg.SOLVER.BASE_LR = args.lr" /home/appuser/trainer/task.py
+RUN sudo sed -i "185 a\    cfg.OUTPUT_DIR = args.output_dir" /home/appuser/trainer/task.py
+RUN sudo sed -i "186 a\    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_file_copy)" /home/appuser/trainer/task.py
+RUN sudo sed -i "182 a\    config_file_copy = args.config_file" /home/appuser/trainer/task.py
+RUN sudo sed -i "183 a\    args.config_file = model_zoo.get_config_file(args.config_file)" /home/appuser/trainer/task.py
+# For new import:
+RUN sudo sed -i "27 a\from detectron2 import model_zoo" /home/appuser/trainer/task.py
+RUN sudo sed -i "21 a\import trainer_utils" /home/appuser/trainer/task.py
+
+ENV PYTHONPATH /home/appuser/trainer
+
+ENTRYPOINT ["python", "-m", "trainer.task"]
@@ -0,0 +1,154 @@
+"""Custom handler for Detectron2 serving."""
+
+import io
+import json
+import os
+from typing import Any, List, Tuple
+
+import cv2
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from google.cloud import storage
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+
+
+def get_bucket_and_blob_name(gcs_filepath: str) -> Tuple[str, str]:
+  """Gets bucket and blob name from gcs path."""
+  # The gcs path is of the form gs://<bucket-name>/<blob-name>
+  gs_suffix = gcs_filepath.split("gs://", 1)[1]
+  return tuple(gs_suffix.split("/", 1))
+
+
+def download_gcs_file(src_file_path: str, dst_file_path: str):
+  """Downloads gcs-file to local folder."""
+  src_bucket_name, src_blob_name = get_bucket_and_blob_name(src_file_path)
+  client = storage.Client()
+  src_bucket = client.get_bucket(src_bucket_name)
+  src_blob = src_bucket.blob(src_blob_name)
+  src_blob.download_to_filename(dst_file_path)
+
+
+class ModelHandler:
+  """Custom model handler for Detectron2."""
+
+  def __init__(self):
+    self.error = None
+    self._batch_size = 0
+    self.initialized = False
+    self.predictor = None
+    self.test_threshold = 0.5
+
+  def initialize(self, context: Any):
+    """Initialize."""
+    print("context.system_properties: ", context.system_properties)
+    print("context.manifest: ", context.manifest)
+    self.manifest = context.manifest
+    properties = context.system_properties
+    # Get threshold from environment variable.
+    # This will be set by customer.
+    self.test_threshold = float(os.environ.get("TEST_THRESHOLD"))
+    print("test_threshold: ", self.test_threshold)
+    # Get model and config file location from environment variables.
+    # These will be set by customer when doing model upload.
+    gcs_model_file = os.environ["MODEL_PTH_FILE"]
+    gcs_config_file = os.environ["CONFIG_YAML_FILE"]
+    print("Copying gcs_model_file: ", gcs_model_file)
+    print("Copying gcs_config_file: ", gcs_config_file)
+    # Copy these files from GCS location to local file.
+    # Note(lavrai): GCSFuse path does not seem to work here for now.
+    model_file = "./model.pth"
+    config_file = "./cfg.yaml"
+    download_gcs_file(src_file_path=gcs_model_file, dst_file_path=model_file)
+    if not os.path.exists(model_file):
+      raise RuntimeError("Missing model_file: %s" % model_file)
+    download_gcs_file(src_file_path=gcs_config_file, dst_file_path=config_file)
+    if not os.path.exists(config_file):
+      raise RuntimeError("Missing config_file: %s" % config_file)
+
+    # Set up config file.
+    cfg = get_cfg()
+    cfg.merge_from_file(config_file)
+    cfg.MODEL.WEIGHTS = model_file
+    cfg.MODEL.DEVICE = (
+        cfg.MODEL.DEVICE + str(properties.get("gpu_id"))
+        if torch.cuda.is_available()
+        else "cpu"
+    )
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.test_threshold
+
+    # Build predictor from config.
+    self.predictor = DefaultPredictor(cfg)
+    self._batch_size = context.system_properties["batch_size"]
+    self.initialized = True
+
+  def preprocess(self, batch: List[Any]) -> List[Any]:
+    """Preprocess raw input and return as list of images."""
+    print("Running pre-processing.")
+    images = []
+    for request in batch:
+      request_data = request.get("data")
+      input_bytes = io.BytesIO(request_data)
+      img = cv2.imdecode(np.fromstring(input_bytes.read(), np.uint8), 1)
+      images.append(img)
+    return images
+
+  def inference(self, model_input: List[Any]) -> List[Any]:
+    """Runs inference."""
+    print("Running model-inference.")
+    return [self.predictor(image) for image in model_input]
+
+  def postprocess(self, inference_result: List[Any]) -> List[Any]:
+    """Post process inference result."""
+    response_list = []
+    print("Num inference_items are:", len(inference_result))
+    for inference_item in inference_result:
+      predictions = inference_item["instances"].to("cpu")
+      print("Predictions are:", predictions)
+      boxes = None
+      if predictions.has("pred_boxes"):
+        boxes = predictions.pred_boxes.tensor.numpy().tolist()
+      scores = None
+      if predictions.has("scores"):
+        scores = predictions.scores.numpy().tolist()
+      classes = None
+      if predictions.has("pred_classes"):
+        classes = predictions.pred_classes.numpy().tolist()
+      masks_rle = None
+      if predictions.has("pred_masks"):
+        # Do run length encoding, else the mask output becomes huge.
+        masks_rle = [
+            mask_util.encode(np.asfortranarray(mask))
+            for mask in predictions.pred_masks
+        ]
+        for rle in masks_rle:
+          rle["counts"] = rle["counts"].decode("utf-8")
+      response = {
+          "classes": classes,
+          "scores": scores,
+          "boxes": boxes,
+          "masks_rle": masks_rle,
+      }
+      response_list.append(json.dumps(response))
+    print("response_list: ", response_list)
+    return response_list
+
+  def handle(self, data: Any, context: Any) -> List[Any]:  # pylint: disable=unused-argument
+    """Runs preprocess, inference, and post-processing."""
+    model_input = self.preprocess(data)
+    model_out = self.inference(model_input)
+    output = self.postprocess(model_out)
+    print("Done handling input.")
+    return output
+
+
+_service = ModelHandler()
+
+
+def handle(data: Any, context: Any) -> List[Any]:
+  if not _service.initialized:
+    _service.initialize(context)
+  if data is None:
+    return None
+  return _service.handle(data, context)
@@ -0,0 +1,97 @@
+"""Detectron2 trainer helper functions."""
+
+import argparse
+from detectron2.data.datasets import register_coco_instances
+
+
+def extend_parser_arguments(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+  """Adds additional model-garden related arguments."""
+  parser.add_argument(
+      "--train_dataset_name",
+      required=False,
+      default="",
+      type=str,
+      help=(
+          "The training dataset name for registration. "
+          "For example: 'balloon_train'."
+      ),
+  )
+  parser.add_argument(
+      "--train_coco_json_file",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the training coco-json format file.",
+  )
+  parser.add_argument(
+      "--train_image_root",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the root folder containing the training images.",
+  )
+  parser.add_argument(
+      "--val_dataset_name",
+      required=False,
+      default="",
+      type=str,
+      help=(
+          "The validation dataset name for registration. "
+          "For example: 'balloon_val'."
+      ),
+  )
+  parser.add_argument(
+      "--val_coco_json_file",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the validation coco-json format file.",
+  )
+  parser.add_argument(
+      "--val_image_root",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the root folder containing the validation images.",
+  )
+  parser.add_argument(
+      "--output_dir",
+      required=True,
+      type=str,
+      help="The path to the output directory.",
+  )
+  # Add hyper-parameter tuning related variables.
+  parser.add_argument(
+      "--lr",
+      type=float,
+      default=0.00025,
+      help="The learning rate to be tuned.",
+  )
+  parser.add_argument(
+      "--hp_eval_task",
+      type=str,
+      choices=["bbox", "segm"],
+      default="bbox",
+      help="The task choice for HP tuning.",
+  )
+  return parser
+
+
+def register_dataset(args: argparse.Namespace):
+  """Register the input dataset in Detectron2 Coco format."""
+  if args.train_dataset_name:
+    register_coco_instances(
+        name=args.train_dataset_name,
+        metadata={},
+        json_file=args.train_coco_json_file,
+        image_root=args.train_image_root,
+    )
+  if args.val_dataset_name:
+    register_coco_instances(
+        name=args.val_dataset_name,
+        metadata={},
+        json_file=args.val_coco_json_file,
+        image_root=args.val_image_root,
+    )