Add requirement for fetch model

vectornguyen76 · Oct 30, 2023 · 4719c45 · 4719c45
1 parent 76637a9
commit 4719c45
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/...el_repository/convert_onnx_to_tensorrt.py → ...gine/model_repository/onnx_to_tensorrt.py b/...el_repository/convert_onnx_to_tensorrt.py → ...gine/model_repository/onnx_to_tensorrt.py
@@ -19,13 +19,13 @@
 def parse_args():
     parser = argparse.ArgumentParser(description="Convert ONNX models to TensorRT")
 
-    parser.add_argument("--device", help="cuda or not", default="cpu")
+    parser.add_argument("--device", help="cuda or not", default="cuda")
 
     # Sample image
     parser.add_argument("--batch_size", type=int, help="data batch size", default=1)
     parser.add_argument("--img_size", help="input size", default=[3, 300, 300])
     parser.add_argument(
-        "--sample_folder_path", help="sample image folder path", default="./test/"
+        "--sample_folder_path", help="sample image folder path", default="./../assets/"
     )
     # parser.add_argument('--sample_image_path', help='sample image path',
     # default='./sample.jpg')

diff --git a/image-search-engine/model_repository/requirements.txt b/image-search-engine/model_repository/requirements.txt
@@ -0,0 +1,45 @@
+charset-normalizer==3.3.1
+cmake==3.27.7
+coloredlogs==15.0.1
+cuda-python==12.3.0
+filelock==3.13.0
+flatbuffers==23.5.26
+humanfriendly==10.0
+idna==3.4
+Jinja2==3.1.2
+lit==17.0.3
+MarkupSafe==2.1.3
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.1
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cublas-cu12==12.3.2.9
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-nvrtc-cu12==12.3.52
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cuda-runtime-cu12==12.3.52
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cudnn-cu12==8.9.4.25
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+onnx==1.15.0
+onnxruntime==1.16.1
+packaging==23.2
+Pillow==10.1.0
+protobuf==4.24.4
+requests==2.31.0
+sympy==1.12
+tensorrt==8.6.1.post1
+tensorrt-bindings==8.6.1
+tensorrt-libs==8.6.1
+torch==1.13.0
+torchaudio==0.13.0
+torchvision==0.14.0
+triton==2.0.0
+typing_extensions==4.8.0
+urllib3==2.0.7
diff --git a/...model_repository/convert_torch_to_onnx.py → ...-engine/model_repository/torch_to_onnx.py b/...model_repository/convert_torch_to_onnx.py → ...-engine/model_repository/torch_to_onnx.py
@@ -14,15 +14,15 @@
 def parse_args():
     parser = argparse.ArgumentParser(description="Convert Pytorch models to ONNX")
 
-    parser.add_argument("--device", help="cuda or not", default="cpu")
+    parser.add_argument("--device", help="cuda or not", default="cuda")
 
     # Sample image
     parser.add_argument(
         "--batch_size", type=int, help="onnx sample batch size", default=1
     )
     parser.add_argument("--img_size", help="image size", default=[3, 300, 300])
     parser.add_argument(
-        "--sample_folder_path", help="sample image folder path", default="./test/"
+        "--sample_folder_path", help="sample image folder path", default="./../assets/"
     )
     # parser.add_argument('--sample_image_path', help='sample image path',
     # default='./sample.jpg')
@@ -183,8 +183,8 @@ def load_image_folder(folder_path, img_size, batch_size):
     onnx.checker.check_model(onnx_model)
 
     # Print a human readable representation of the graph
-    with open("OnnxShape.txt", "w") as f:
-        f.write(f"{onnx.helper.printable_graph(onnx_model.graph)}")
+    # with open("OnnxShape.txt", "w") as f:
+    #     f.write(f"{onnx.helper.printable_graph(onnx_model.graph)}")
 
     # Comparision output of onnx and output of Pytorch model
     # Pytorch results