feat: update main scripts with TracingAdapter and hydra config

.gitignore

@@ -176,4 +176,6 @@ res.png
 *.ep
 output/*
 *.png
-docs/_build
+docs/_build
+!docs/src/**/*.png
+outputs/

projects/dino_dinov2/modeling/exportable/dino_transformer.py

@@ -305,17 +305,6 @@ def gen_encoder_output_proposals(self, memory, memory_padding_mask, spatial_shap
         N, S, C = memory.shape
         proposals = []
         _cur = 0
-        # temporary
-        # REFACTOR: we can remove these static checks if we never cast the spatial_shapes to a tensor
-        # spatial_shapes: List[Tuple[int, int]] = spatial_shapes.tolist()
-        # torch._check(spatial_shapes[0][0] == 150)
-        # torch._check(spatial_shapes[0][1] == 100)
-        # torch._check(spatial_shapes[1][0] == 75)
-        # torch._check(spatial_shapes[1][1] == 50)
-        # torch._check(spatial_shapes[2][0] == 37)
-        # torch._check(spatial_shapes[2][1] == 25)
-        # torch._check(spatial_shapes[3][0] == 19)
-        # torch._check(spatial_shapes[3][1] == 13)
         for lvl, (H, W) in enumerate(spatial_shapes):
             mask_flatten_ = memory_padding_mask[:, _cur : (_cur + H * W)].view(
                 N, H, W, 1
@@ -375,26 +364,8 @@ def get_reference_points(spatial_shapes: List[Tuple[int, int]], valid_ratios: to
             Tensor: reference points used in decoder, has \
                 shape (bs, num_keys, num_levels, 2).
         """
-        # spatial_shapes: list[tuple[int, int]] =  spatial_shapes.tolist()
-        # [[150, 100], [75, 50], [37, 25], [19, 13]]
-        # temporary
-        # torch._check(spatial_shapes[0][0] == 150)
-        # torch._check(spatial_shapes[0][1] == 100)
-        # torch._check(spatial_shapes[1][0] == 75)
-        # torch._check(spatial_shapes[1][1] == 50)
-        # torch._check(spatial_shapes[2][0] == 37)
-        # torch._check(spatial_shapes[2][1] == 25)
-        # torch._check(spatial_shapes[3][0] == 19)
-        # torch._check(spatial_shapes[3][1] == 13)
         reference_points_list = []
         for lvl, (H, W) in enumerate(spatial_shapes):
-            #  TODO  check this 0.5
-            # H, W = H.item(), W.item()
-            # torch._check(0 <= H)
-            # torch._check(H <= 200)
-            # torch._check(0 <= W)
-            # torch._check(W <= 200)
-
             ref_y, ref_x = torch.meshgrid(
                 torch.linspace(0.5, H - 0.5, H, dtype=torch.float32, device=device),
                 torch.linspace(0.5, W - 0.5, W, dtype=torch.float32, device=device),

pyproject.toml

@@ -4,14 +4,16 @@ version = "0.1.0"
 description = ""
 authors = ["None"]
 readme = "README.md"
+packages = [
+    { include = "src" }
+]
 
 [tool.poetry.dependencies]
 python = "3.10"
-torch = {version = "^2.6.0.dev20241013+cu124", allow-prereleases = true, source = "torch-cu124"}
-torchvision = {version = "^0.20.0.dev20241013+cu124", allow-prereleases = true, source = "torch-cu124"}
-torch-tensorrt = {version = "^2.6.0.dev20241013+cu124", allow-prereleases = true, source = "torch-cu124"}
+torch = {version = "^2.6.0.dev20241027+cu124", allow-prereleases = true, source = "torch-cu124"}
+torchvision = {version = "^0.20.0.dev20241027+cu124", allow-prereleases = true, source = "torch-cu124"}
+torch-tensorrt = {version = "^2.6.0.dev20241027+cu124", allow-prereleases = true, source = "torch-cu124"}
 tensorrt = "10.3.0"
-nvidia-cuda-runtime-cu12 = "12.4.127"
 nvidia-modelopt = {extras = ["all"], version = "^0.17.0", source = "pypi-nvidia"}
 onnxscript = "^0.1.0.dev20241017"
 
@@ -41,3 +43,6 @@ livereload = "^2.7.0"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.isort]
+profile = "black"

scripts/_modelopt_quantize_onnx.py

@@ -1,7 +1,7 @@
 import torch_tensorrt
 import torch
 import modelopt.torch.quantization as mtq
-from src.utils.quantization import load_input_fixed, ModelWrapper, load_model, unflatten_repr, plot_predictions
+from src.utils._quantization import load_input_fixed, ModelWrapper, load_model, unflatten_repr, plot_predictions
 # from torchao.utils import unwrap_tensor_subclass
 from modelopt.torch.quantization.utils import export_torch_mode
 

scripts/_modelopt_quantize_pt2e.py

@@ -1,7 +1,7 @@
 import torch_tensorrt
 import torch
 import modelopt.torch.quantization as mtq
-from src.utils.quantization import load_input_fixed, ModelWrapper, load_model, unflatten_repr, plot_predictions
+from src.utils._quantization import load_input_fixed, ModelWrapper, load_model, unflatten_repr, plot_predictions
 # from torchao.utils import unwrap_tensor_subclass
 from modelopt.torch.quantization.utils import export_torch_mode
 

scripts/_quantize_cpu.py

@@ -1,7 +1,7 @@
 import torch
 import torchao
 import detectron2.structures
-from src.utils.quantization import (
+from src.utils._quantization import (
     load_input,
     load_model,
     register_DINO_output_types,

scripts/_quantize_trt.py

@@ -3,7 +3,7 @@
 from torch.ao.quantization.quantizer.xnnpack_quantizer import XNNPACKQuantizer, get_symmetric_quantization_config
 from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
 import detectron2.structures
-from src.utils.quantization import (
+from src.utils._quantization import (
     load_input,
     load_model,
     register_DINO_output_types,

scripts/_torch_compile.py

@@ -0,0 +1,8 @@
+import torch
+def f(x):
+    return x + 1
+
+
+x = torch.randn(1, 3, 224, 224)
+m = torch.compile(f)
+print(m(x)) 

scripts/_torchao_quantize.py

@@ -1,7 +1,7 @@
 import torch
 # import torch_tensorrt
 # import modelopt.torch.quantization as mtq
-from src.utils.quantization import load_input_fixed, ModelWrapper, load_model
+from src.utils._quantization import load_input_fixed, ModelWrapper, load_model
 from torchao.utils import unwrap_tensor_subclass
 import pickle
 import torchao

scripts/benchmark_gpu.py

@@ -3,13 +3,12 @@
 from typing import Optional
 from functools import partial
 import contextlib
-from src.utils.quantization import (
+from src.utils import (
     load_input_fixed,
-    unflatten_repr,
     plot_predictions,
-    ModelWrapper,
+    TracingAdapter,
 )
-from src.utils.quantization import load_model as _load_model
+from src.utils import load_model as _load_model
 from statistics import stdev, mean
 import torch_tensorrt
 import logging
@@ -19,6 +18,7 @@
 
 detrex.layers.multi_scale_deform_attn._ENABLE_CUDA_MSDA = False
 
+
 def setup_parser():
     DEFAULT_IMG = Path("artifacts/idea_raw.jpg")
     parser = argparse.ArgumentParser()
@@ -27,7 +27,9 @@ def setup_parser():
     parser.add_argument("--n_warmup", type=int, default=10)
     parser.add_argument("--n_iter", type=int, default=10)
     parser.add_argument("--output", type=Path, default=None)
-    parser.add_argument("--amp_dtype", type=str, default=None, choices=["fp16", "bf16", None])
+    parser.add_argument(
+        "--amp_dtype", type=str, default=None, choices=["fp16", "bf16", None]
+    )
     return parser
 
 
@@ -43,11 +45,8 @@ def load_model(model_path: Path):
         model = torch.export.load(model_path).module()
     elif model_path.suffix == ".pth":
         height, width = 512, 512
-        model = ModelWrapper(
-            net=_load_model().cuda(),
-            height=height,
-            width=width,
-        )
+        model = _load_model().cuda()
+        model = TracingAdapter(model, *load_input_fixed(height=height, width=width))
     else:
         raise ValueError(f"Unsupported model format: {model_path.suffix}")
 

scripts/config/export_tensorrt/dinov2.yaml

@@ -0,0 +1,15 @@
+image:
+  height: 512
+  width: 512
+  path: "artifacts/idea_raw.jpg"
+amp_dtype: "fp32"
+trt:
+  enabled_precisions:
+    - "fp32"
+model:
+  config: "projects/dino_dinov2/configs/models/dino_dinov2.py"
+  ckpt_path: "artifacts/model_final.pth"
+  opts:
+    - "model.backbone.net.img_size=[512, 512]"
+    - "model.backbone.net.dynamic_img_size=False"
+    - "model.backbone.net.dynamic_img_pad=False"

scripts/config/export_tensorrt/vit.yaml

@@ -0,0 +1,17 @@
+model:
+  config: "detrex/detectron2/projects/ViTDet/configs/COCO/cascade_mask_rcnn_vitdet_b_100ep.py"
+  ckpt_path: "artifacts/model_final_435fa9.pkl"
+  opts:
+    - "model.backbone.net.window_block_indexes=[]"
+image:
+  height: 1024
+  width: 1024
+  path: "artifacts/idea_raw.jpg"
+trt:
+  enabled_precisions: 
+    - "fp32"
+  require_full_compilation: False
+  enable_experimental_decompositions: True
+  min_block_size: 1
+  use_fast_partitioner: True # doesn't make any difference in supported nodes
+amp_dtype: "fp32"