add ci unit tests

.github/workflows/tests.yml

@@ -0,0 +1,80 @@
+name: Unit Tests
+
+on:
+  pull_request:
+    branches: [main]
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: unit-tests-${{ github.ref }}
+  cancel-in-progress: true
+
+env:
+  IMAGE: sam3-tensorrt:infer
+
+jobs:
+  unit-tests:
+    runs-on: [self-hosted, linux, gpu]   # must match --labels used in config.sh
+    timeout-minutes: 90
+
+    if: >-
+      github.event_name != 'pull_request' ||
+      github.event.pull_request.head.repo.full_name == github.repository
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          clean: false
+
+      - name: Build docker image
+        working-directory: docker
+        run: ./build-trt.bash
+
+      - name: Run unit tests in container
+        env:
+          MODELS_DIR: /home/jason/Projects/sam3-ws/SAM3-TensorRT/models
+          TEST_SCRIPT: |
+            set -uo pipefail
+            rc=0
+
+            if ! ls "$HOME"/models/*.engine >/dev/null 2>&1; then
+              echo "::error::No .engine files found in the mounted models dir"
+              exit 1
+            fi
+            nvidia-smi || true
+
+            echo "::group::C++ build and tests"
+            ( cd "$HOME/sam3trt" \
+              && cmake -S . -B build -DSAM3TRT_BUILD_TESTS=ON -DSAM3TRT_BUILD_ROS=OFF \
+              && cmake --build build -j"$(nproc)" \
+              && ctest --test-dir build --output-on-failure -E Timing ) || rc=1
+            echo "::endgroup::"
+
+            echo "::group::Python tests"
+            ( cd "$HOME" \
+              && pip install --user --break-system-packages --no-cache-dir ./sam3trtpy \
+              && cd sam3trtpy \
+              && python3 -m unittest discover -s tests -v ) || rc=1
+            echo "::endgroup::"
+
+            exit $rc
+        run: |
+          U="$(whoami)"
+          # Same mounts as docker/run-trt.bash, minus the interactive/X11 parts.
+          docker run --rm \
+            --gpus all \
+            --ipc=host \
+            --user "$(id -u):$(id -g)" \
+            -e HOME="/home/$U" \
+            -e TEST_SCRIPT \
+            -v "$MODELS_DIR:/home/$U/models" \
+            -v "$GITHUB_WORKSPACE/sam3trt:/home/$U/sam3trt" \
+            -v "$GITHUB_WORKSPACE/sam3trtpy:/home/$U/sam3trtpy" \
+            -w "/home/$U" \
+            "$IMAGE" \
+            bash -c "$TEST_SCRIPT"

docker/Dockerfile.trt

@@ -35,6 +35,9 @@ WORKDIR /home/$USER
 
 RUN sudo apt install -y python3-pip
 RUN pip3 install --no-deps opencv-python-headless ultralytics
+RUN pip install --no-cache-dir ultralytics==8.4.121
+ENV LD_LIBRARY_PATH=/opt/hpcx/ucx/lib:/opt/hpcx/ucc/lib:${LD_LIBRARY_PATH}
+
 
 RUN sudo apt update \
  && sudo apt install -y \

scripts/export_ultralytics_image_encoder.py

@@ -80,12 +80,6 @@ def patch_layernorm(model):
 def trace_and_export_image_encoder(model : torch.nn.Module, input : Any, engine_path : str, fp16 : bool) -> Tuple[torch.Tensor]: 
     wrapper = ImageEncoderWrapper(model).to(DEVICE).eval()
 
-    predictor.setup_source(input)
-
-    for batch in predictor.dataset:
-        input = predictor.preprocess(batch[1])
-        break
-
     print("[ImageEncoderExport] Tracing Image Encoder Model")
     torch_output = wrapper(input) # warm up
     exp_program = torch.export.export(wrapper, (input,), strict=False)
@@ -223,16 +217,17 @@ def export_and_verify_image_encoder(predictor : SAM3SemanticPredictor, fp16 : bo
 
     precision = "fp16" if fp16 else "fp32"
     engine_path = os.path.join(os.environ["HOME"], "models", f"image_encoder_{precision}.engine")
-    #_ = trace_and_export_image_encoder(
-    #    predictor.model,
-    #    img,
-    #    engine_path,
-    #    fp16
-    #)
 
     for batch in predictor.dataset:
         im = predictor.preprocess(batch[1])
         break
+    _ = trace_and_export_image_encoder(
+        predictor.model,
+        im,
+        engine_path,
+        fp16
+    )
+
     _verify_engine(engine_path, im, torch_output, min_cos=0.999 if fp16 else 0.9999)
 
     return torch_output
@@ -255,15 +250,16 @@ def export_and_verify_image_encoder(predictor : SAM3SemanticPredictor, fp16 : bo
 
     precision = "fp16" if args.fp16 else "fp32"
     engine_path = os.path.join(os.environ["HOME"], "models", f"image_encoder_{precision}.engine")
+
+    for batch in predictor.dataset:
+        im = predictor.preprocess(batch[1])
+        break
     _ = trace_and_export_image_encoder(
         predictor.model,
-        img,
+        im,
         engine_path,
         args.fp16
     )
 
-    for batch in predictor.dataset:
-        im = predictor.preprocess(batch[1])
-        break
     _verify_engine(engine_path, im, torch_output, min_cos=0.999 if args.fp16 else 0.9999) 
 

scripts/export_ultralytics_mask_decoder.py

@@ -238,7 +238,7 @@ def trace_and_export_mask_deocder(model : torch.nn.Module, input : Any, engine_p
     print("[MaskDecoderExport] Patching Mask Decoder Model")
     _, _, spatial_shapes, _ = fusion_wrapper(*input)
     H, W = int(spatial_shapes[0, 0]), int(spatial_shapes[0, 1])
-    patch_rpb(predictor.model, H, W)
+    patch_rpb(model, H, W)
 
     torch_output = wrapper(*input)
 
@@ -256,9 +256,10 @@ def trace_and_export_mask_deocder(model : torch.nn.Module, input : Any, engine_p
             offload_module_to_cpu=True,
             device=torch_tensorrt.Device("cuda:0"),
     )
+    model.to(DEVICE)  # offload_module_to_cpu leaves the weights on the CPU
 
-    #with open(engine_path, "wb") as f:
-    #    f.write(engine_bytes)
+    with open(engine_path, "wb") as f:
+        f.write(engine_bytes)
 
     return torch_output 
 
@@ -376,12 +377,12 @@ def export_and_verify_mask_decoder(predictor : SAM3SemanticPredictor, input : Tu
 
     precision =  "fp16" if fp16 else "fp32"
     engine_path = os.path.join(os.environ["HOME"], "models", f"mask_decoder_{precision}.engine")
-    #trace_and_export_mask_deocder(
-    #    predictor.model, 
-    #    input,
-    #    engine_path,
-    #    fp16
-    #)
+    trace_and_export_mask_deocder(
+        predictor.model, 
+        input,
+        engine_path,
+        fp16
+    )
     predictor.set_image(img)
     ref = predictor(text=captions)[0] 
     for batch in predictor.dataset:

scripts/export_ultralytics_text_encoder.py

@@ -79,7 +79,7 @@ def trace_and_export_text_encoder(model : torch.nn.Module, input : Any, engine_p
         exp_program,
         arg_inputs=[input[0].to(DEVICE), input[1].to(DEVICE)],
         optimization_level=5,
-        use_explicit_typeing=True,
+        use_explicit_typing=True,
         device=torch_tensorrt.Device(f"cuda:0"),
     )
 
@@ -184,7 +184,7 @@ def export_and_verify_text_encoder(predictor : SAM3SemanticPredictor, fp16 : boo
 
     precision = "fp16" if fp16 else "fp32"
     engine_path = os.path.join(os.environ["HOME"], "models", f"text_encoder_{precision}.engine")
-    #trace_and_export_text_encoder(predictor.model, (input_ids, attention_mask), engine_path, args.fp16)
+    trace_and_export_text_encoder(predictor.model, (input_ids, attention_mask), engine_path, fp16)
 
     _verify_engine(engine_path, (input_ids, attention_mask), (torch_output[0], torch_output[2]), 0.999 if fp16 else 0.9999)