diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..d4ccf7b --- /dev/null +++ b/.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" diff --git a/docker/Dockerfile.trt b/docker/Dockerfile.trt index bc7490d..45460b0 100644 --- a/docker/Dockerfile.trt +++ b/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 \ diff --git a/scripts/export_ultralytics_image_encoder.py b/scripts/export_ultralytics_image_encoder.py index 26a7f5c..2488f92 100644 --- a/scripts/export_ultralytics_image_encoder.py +++ b/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) diff --git a/scripts/export_ultralytics_mask_decoder.py b/scripts/export_ultralytics_mask_decoder.py index 606c4ff..02836f7 100644 --- a/scripts/export_ultralytics_mask_decoder.py +++ b/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: diff --git a/scripts/export_ultralytics_text_encoder.py b/scripts/export_ultralytics_text_encoder.py index ecc6dd7..b7a44b1 100644 --- a/scripts/export_ultralytics_text_encoder.py +++ b/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)