GPU in docker

ispras · Dec 14, 2023 · 09632a7 · 09632a7
1 parent f3ec0e5
commit 09632a7
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -6,6 +6,8 @@ ENV RESOURCES_PATH "/dedoc_root/resources"
 
 ADD requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt
+RUN pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
+
 
 RUN mkdir /dedoc_root
 RUN mkdir /dedoc_root/dedoc

diff --git a/dedoc/config.py b/dedoc/config.py
@@ -24,7 +24,7 @@
 
     # --------------------------------------------GPU SETTINGS----------------------------------------------------------
     # set gpu in XGBoost and torch models
-    on_gpu=False,
+    on_gpu=True,
 
     # ---------------------------------------------API SETTINGS---------------------------------------------------------
     # max file size in bytes

diff --git a/..._reader/pdf_image_reader/columns_orientation_classifier/columns_orientation_classifier.py b/..._reader/pdf_image_reader/columns_orientation_classifier/columns_orientation_classifier.py
@@ -38,8 +38,9 @@ def net(self) -> ClassificationModelTorch:
                 self._load_weights(net)
             else:
                 net = ClassificationModelTorch(None)
-            net.to(self.device)
             self._nets[self.checkpoint_path] = net
+        if self.device != next(self._nets[self.checkpoint_path].parameters()).device:
+            self._nets[self.checkpoint_path].to(self.device)
         return self._nets[self.checkpoint_path]
 
     @staticmethod

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -12,7 +12,8 @@ services:
       - 1231:1231
     environment:
       DOCREADER_PORT: 1231
-
+      NVIDIA_VISIBLE_DEVICES: 0
+    runtime: nvidia
 
   test:
     depends_on:
@@ -26,6 +27,8 @@ services:
       DOCREADER_PORT: 1231
       is_test: $test
       PYTHONPATH: $PYTHONPATH:/dedoc_root/tests:/dedoc_root
+      NVIDIA_VISIBLE_DEVICES: 0
+    runtime: nvidia
     command:
         bash dedoc_root/tests/run_tests_in_docker.sh