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