AOTInductor Resnet CPP example

cpp/CMakeLists.txt

@@ -30,7 +30,7 @@ find_package(folly REQUIRED)
 find_package(fmt REQUIRED)
 find_package(gflags REQUIRED)
 find_package(Torch REQUIRED)
-find_package(yaml-cpp REQUIRED NO_CMAKE_PATH)
+
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
 
 include_directories(${TORCH_INCLUDE_DIRS})

cpp/build.sh

@@ -155,37 +155,6 @@ function install_yaml_cpp() {
   cd "$BWD" || exit
 }
 
-function install_sentencepiece() {
-  SENTENCEPIECE_SRC_DIR=$BASE_DIR/third-party/sentencepiece
-  SENTENCEPIECE_BUILD_DIR=$DEPS_DIR/sentencepiece-build
-
-  if [ ! -d "$SENTENCEPIECE_SRC_DIR" ] ; then
-    echo -e "${COLOR_GREEN}[ INFO ] Cloning sentencepiece repo ${COLOR_OFF}"
-    git clone https://github.com/google/sentencepiece.git "$SENTENCEPIECE_SRC_DIR"
-    cd $SENTENCEPIECE_SRC_DIR
-    git checkout tags/v0.1.99
-  fi
-
-  if [ ! -d "$SENTENCEPIECE_BUILD_DIR" ] ; then
-    echo -e "${COLOR_GREEN}[ INFO ] Building sentencepiece ${COLOR_OFF}"
-
-    mkdir $SENTENCEPIECE_BUILD_DIR
-    cd $SENTENCEPIECE_BUILD_DIR
-    cmake $SENTENCEPIECE_SRC_DIR
-    make -i $(nproc)
-    if [ "$PLATFORM" = "Linux" ]; then
-      sudo make install
-      sudo ldconfig -v
-    elif [ "$PLATFORM" = "Mac" ]; then
-      make install
-    fi
-
-    echo -e "${COLOR_GREEN}[ INFO ] sentencepiece is installed ${COLOR_OFF}"
-  fi
-
-  cd "$BWD" || exit
-}
-
 function build_llama_cpp() {
   BWD=$(pwd)
   LLAMA_CPP_SRC_DIR=$BASE_DIR/third-party/llama.cpp
@@ -208,14 +177,23 @@ function prepare_test_files() {
   if [ ! -f "${EX_DIR}/babyllama/babyllama_handler/stories15M.bin" ]; then
     wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin -O "${EX_DIR}/babyllama/babyllama_handler/stories15M.bin"
   fi
-  if [ ! -f "${EX_DIR}/aot_inductor/llama_handler/stories15M.so" ]; then
-    local HANDLER_DIR=${EX_DIR}/aot_inductor/llama_handler/
-    if [ ! -f "${HANDLER_DIR}/stories15M.pt" ]; then
-      wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.pt?download=true -O "${HANDLER_DIR}/stories15M.pt"
+  # PT2.2 torch.expport does not support Mac
+  if [ "$PLATFORM" = "Linux" ]; then
+    if [ ! -f "${EX_DIR}/aot_inductor/llama_handler/stories15M.so" ]; then
+      local HANDLER_DIR=${EX_DIR}/aot_inductor/llama_handler/
+      if [ ! -f "${HANDLER_DIR}/stories15M.pt" ]; then
+        wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.pt?download=true -O "${HANDLER_DIR}/stories15M.pt"
+      fi
+      local LLAMA_SO_DIR=${BASE_DIR}/third-party/llama2.so/
+      PYTHONPATH=${LLAMA_SO_DIR}:${PYTHONPATH} python ${BASE_DIR}/../examples/cpp/aot_inductor/llama2/compile.py --checkpoint ${HANDLER_DIR}/stories15M.pt ${HANDLER_DIR}/stories15M.so
+    fi
+    if [ ! -f "${EX_DIR}/aot_inductor/resnet_handler/resne50_pt2.so" ]; then
+      local HANDLER_DIR=${EX_DIR}/aot_inductor/resnet_handler/
+      cd ${HANDLER_DIR}
+      python ${BASE_DIR}/../examples/cpp/aot_inductor/resnet/resnet50_torch_export.py
     fi
-    local LLAMA_SO_DIR=${BASE_DIR}/third-party/llama2.so/
-    PYTHONPATH=${LLAMA_SO_DIR}:${PYTHONPATH} python ${BASE_DIR}/../examples/cpp/aot_inductor/llama2/compile.py --checkpoint ${HANDLER_DIR}/stories15M.pt ${HANDLER_DIR}/stories15M.so
   fi
+  cd "$BWD" || exit
 }
 
 function build() {
@@ -401,7 +379,6 @@ install_folly
 install_kineto
 install_libtorch
 install_yaml_cpp
-install_sentencepiece
 build_llama_cpp
 prepare_test_files
 build

cpp/src/examples/CMakeLists.txt

@@ -1,8 +1,14 @@
 
 add_subdirectory("../../../examples/cpp/babyllama/" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/babyllama/babyllama_handler/")
 
-add_subdirectory("../../../examples/cpp/aot_inductor/llama2/" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/aot_inductor/llama_handler/")
-
 add_subdirectory("../../../examples/cpp/llamacpp/" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/llamacpp/llamacpp_handler/")
 
 add_subdirectory("../../../examples/cpp/mnist/" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/mnist/mnist_handler/")
+
+
+# PT2.2 torch.expport does not support Mac
+if(CMAKE_SYSTEM_NAME MATCHES "Linux")
+  add_subdirectory("../../../examples/cpp/aot_inductor/llama2/" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/aot_inductor/llama_handler/")
+
+  add_subdirectory("../../../examples/cpp/aot_inductor/resnet" "${CMAKE_CURRENT_BINARY_DIR}/../../test/resources/examples/aot_inductor/resnet_handler/")
+endif()

cpp/src/utils/CMakeLists.txt

@@ -12,7 +12,12 @@ list(APPEND TS_UTILS_SOURCE_FILES ${TS_UTILS_SRC_DIR}/metrics/registry.cc)
 add_library(ts_utils SHARED ${TS_UTILS_SOURCE_FILES})
 target_include_directories(ts_utils PUBLIC ${TS_UTILS_SRC_DIR})
 target_include_directories(ts_utils PRIVATE ${Boost_INCLUDE_DIRS})
-target_link_libraries(ts_utils ${FOLLY_LIBRARIES} ${CMAKE_DL_LIBS} ${Boost_LIBRARIES} yaml-cpp::yaml-cpp)
+if(CMAKE_SYSTEM_NAME MATCHES "Darwin")
+  target_link_libraries(ts_utils ${FOLLY_LIBRARIES} ${CMAKE_DL_LIBS} ${Boost_LIBRARIES} yaml-cpp::yaml-cpp)
+else()
+  target_link_libraries(ts_utils ${FOLLY_LIBRARIES} ${CMAKE_DL_LIBS} ${Boost_LIBRARIES} yaml-cpp)
+endif()
+
 install(TARGETS ts_utils DESTINATION ${torchserve_cpp_SOURCE_DIR}/_build/libs)
 
 list(APPEND FOO_SOURCE_FILES ${TS_UTILS_SRC_DIR}/ifoo.hh)

cpp/test/examples/examples_test.cc

@@ -59,3 +59,24 @@ TEST_F(ModelPredictTest, TestLoadPredictLlamaCppHandler) {
           base_dir + "llamacpp_handler", "llamacpp", -1, "", "", 1, false),
       base_dir + "llamacpp_handler", base_dir + "prompt.txt", "llm_ts", 200);
 }
+
+TEST_F(ModelPredictTest, TestLoadPredictAotInductorResnetHandler) {
+  std::string base_dir = "_build/test/resources/examples/aot_inductor/";
+  std::string file1 = base_dir + "resnet_handler/resnet50_pt2.so";
+
+  std::ifstream f1(file1);
+
+  if (!f1.good())
+    GTEST_SKIP() << "Skipping TestLoadPredictAotInductorResnetHandler because "
+                    "of missing files: "
+                 << file1;
+
+  this->LoadPredict(
+    std::make_shared<torchserve::LoadModelRequest>(
+      base_dir + "resnet_handler", "resnet50_aot",
+      torch::cuda::is_available() ? 0 : -1, "", "", 1, false),
+      base_dir + "resnet_handler",
+      base_dir + "resnet_handler/0_png.pt",
+      "resnet_ts",
+      200);
+}

cpp/test/resources/examples/aot_inductor/resnet_handler/MAR-INF/MANIFEST.json

@@ -0,0 +1,11 @@
+{
+  "createdOn": "14/02/2024 05:57:14",
+  "runtime": "LSP",
+  "model": {
+    "modelName": "resnetcppaot",
+    "handler": "libresnet_handler:ResnetCppHandler",
+    "modelVersion": "1.0",
+    "configFile": "model-config.yaml"
+  },
+  "archiverVersion": "0.9.0"
+}

cpp/test/resources/examples/aot_inductor/resnet_handler/model-config.yaml

@@ -0,0 +1,7 @@
+minWorkers: 1
+maxWorkers: 1
+batchSize: 2
+
+handler:
+  model_so_path: "resnet50_pt2.so"
+  mapping: "index_to_name.json"

examples/cpp/aot_inductor/resnet/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_library(resnet_handler SHARED src/resnet_handler.cc)
+target_link_libraries(resnet_handler PRIVATE ts_backends_core ts_utils ${TORCH_LIBRARIES})

examples/cpp/aot_inductor/resnet/README.md

@@ -0,0 +1,61 @@
+This example uses AOTInductor to compile the Resnet50 into an so file which is then executed using libtorch.
+The handler C++ source code for this examples can be found [here](src).
+
+### Setup
+1. Follow the instructions in [README.md](../../../../cpp/README.md) to build the TorchServe C++ backend.
+
+```
+cd serve/cpp
+./builld.sh
+```
+
+The build script will create the necessary artifact for this example.
+To recreate these by hand you can follow the prepare_test_files function of the [build.sh](../../../../cpp/build.sh) script.
+We will need the handler .so file as well as the resnet50_pt2.so file containing the model and weights.
+
+2. Create a [model-config.yaml](model-config.yaml)
+
+```yaml
+minWorkers: 1
+maxWorkers: 1
+batchSize: 2
+
+handler:
+  model_so_path: "resnet50_pt2.so"
+  mapping: "index_to_name.json"
+```
+
+### Generate Model Artifacts Folder
+
+```bash
+torch-model-archiver --model-name resnetcppaot --version 1.0 --handler ../../../../cpp/_build/test/resources/examples/aot_inductor/resnet_handler/libresnet_handler:ResnetCppHandler --runtime LSP --extra-files index_to_name.json,../../../../cpp/_build/test/resources/examples/aot_inductor/resnet_handler/resnet50_pt2.so --config-file model-config.yaml --archive-format no-archive
+```
+
+Create model store directory and move the folder `resnetcppaot`
+
+```
+mkdir model_store
+mv resnetcppaot model_store/
+```
+
+### Inference
+
+Start torchserve using the following command
+
+```
+torchserve --ncs --model-store model_store/ --models resnetcppaot
+```
+
+
+Infer the model using the following command
+
+```
+curl http://localhost:8080/predictions/resnetcppaot -T ../../../../cpp/test/resources/examples/aot_inductor/resnet_handler/0_png.pt
+{
+  "lens_cap": 0.0022578993812203407,
+  "lynx": 0.0032067005522549152,
+  "Egyptian_cat": 0.046274684369564056,
+  "tiger_cat": 0.13740436732769012,
+  "tabby": 0.2724998891353607
+}
+```

examples/cpp/aot_inductor/resnet/model-config.yaml

@@ -0,0 +1,7 @@
+minWorkers: 1
+maxWorkers: 1
+batchSize: 2
+
+handler:
+  model_so_path: "resnet50_pt2.so"
+  mapping: "index_to_name.json"

examples/cpp/aot_inductor/resnet/resnet50_torch_export.py

@@ -0,0 +1,40 @@
+import os
+
+import torch
+from torchvision.models import ResNet50_Weights, resnet50
+
+torch.set_float32_matmul_precision("high")
+
+MAX_BATCH_SIZE = 15
+
+model = resnet50(weights=ResNet50_Weights.DEFAULT)
+model.eval()
+
+with torch.no_grad():
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+        # The max batch size is less than 16. The following setting can only work in PT2.3.
+        # We need to turn off the below optimizations to support batch_size = 16,
+        # which is treated like a special case
+        # https://github.com/pytorch/pytorch/pull/116152
+        # torch.backends.mkldnn.set_flags(False)
+        # torch.backends.nnpack.set_flags(False)
+
+    model = model.to(device=device)
+    example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
+
+    batch_dim = torch.export.Dim("batch", min=1, max=MAX_BATCH_SIZE)
+    torch._C._GLIBCXX_USE_CXX11_ABI = True
+    so_path = torch._export.aot_compile(
+        model,
+        example_inputs,
+        # Specify the first dimension of the input x as dynamic
+        dynamic_shapes={"x": {0: batch_dim}},
+        # Specify the generated shared library path
+        options={
+            "aot_inductor.output_path": os.path.join(os.getcwd(), "resnet50_pt2.so"),
+            "max_autotune": True,
+        },
+    )