make docker-files runnable on GPUs

models/DockerfileCUDA

@@ -0,0 +1,12 @@
+FROM pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime
+
+WORKDIR /usr/src/app
+
+RUN apt-get update 
+RUN apt-get install libgl1 -y
+
+COPY requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+

models/docker-compose.cuda.yml

@@ -0,0 +1,26 @@
+version: "3"
+services:
+  model-cuda:
+    build: 
+      context: .
+      dockerfile: DockerfileCUDA
+    tty: true
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+    volumes:
+      - ./logs/:/usr/src/app/logs:z
+      - ./data/:/usr/src/app/data:z
+      - ./configs:/usr/src/app/configs:z
+    command: sh -c "python train.py logger=tensorboard trainer.gpus=1"
+  tensorboard:
+    image: tensorflow/tensorflow
+    volumes:
+      - ./logs:/logs:z
+    command: sh -c "tensorboard --logdir /logs --bind_all"
+    ports:
+      - 8008:6006

models/docker-compose.yml

@@ -1,7 +1,10 @@
 version: "3"
 services:
   model:
-    build: .
+    build: 
+      context: .
+      dockerfile: Dockerfile
+    tty: true
     volumes:
       - ./logs/:/usr/src/app/logs:z
       - ./data/:/usr/src/app/data:z
@@ -13,4 +16,4 @@ services:
       - ./logs:/logs:z
     command: sh -c "tensorboard --logdir /logs --bind_all"
     ports:
-      - 6006:6006
+      - 8008:6006