Docker, docker-compose에 원하는 GPU만 할당하기

현재 내 서버에는 GPU가 2개 있는데 학습하고 있는 코드는 multi-gpu 코드가 아니다. 그러므로 2개의 GPU에 각각 코드를 돌린다면 2번의 실험을 동시에 할 수 있다. 그래서 docker-compose를 이용해 특정 GPU만 할당하는 방법을 알아보았다.

 

0. GPU 확인

nvidia-smi

output:

1070 2개가 보인다. GPU index는 0부터 시작하여 2개니깐 1에서 끝난다.

 

1. docker

docker run -e NVIDIA_VISIBLE_DEVICES=1 hsh0322/las-pytorch python3 train.py

위와 같이 -e NVIDIA_VISIBLE_DEVICES=(원하는 GPU number)로 옵션을 넣어주면 된다.

 

2. docker-compose

version: '3'

services:
  train:
    image: hsh0322/las-pytorch
    environment:
      - NVIDIA_VISIBLE_DEVICES=1
    command:
       python3 train.py

docker-compose.yaml에 environment라는 옵션에 같은 내용을 적어준다.

주의할 점은 docker-compose.yaml의 service 이름이 같다면 docker-compose는 같은 service로 인식하고 container를 recreate한다. 그러므로 다른 GPU로 돌릴 때는 두개의 docker-compose.yaml과 다른 이름의 service가 필요하다.

 

3. docker-compose.yaml 예시

version: '3'

services:
  train0:
    image: hsh0322/las-pytorch
    volumes:
      - ./:/opt/project
    environment:
      - NVIDIA_VISIBLE_DEVICES=0
    command:
       python3 train_timit.py config/las_timit_mfcc_spikegram_150_config.yaml

  tensorboard0:
    image: hsh0322/las-pytorch
    volumes:
      - ./:/opt/project
    ports:
      - "6006:6006"
    command: tensorboard --logdir=log/las_timit --port=6006 --host=0.0.0.0

  train1:
    image: hsh0322/las-pytorch
    volumes:
      - ./:/opt/project
    environment:
      - NVIDIA_VISIBLE_DEVICES=1
    command:
      python3 train_timit.py config/las_timit_mfcc_spikegram_150_config1.yaml

  tensorboard1:
    image: hsh0322/las-pytorch
    volumes:
      - ./:/opt/project
    ports:
      - "6005:6005"
    command: tensorboard --logdir=log/las_timit1 --port=6005 --host=0.0.0.0

위의 예제는 필자가 현재 사용하고 있는 docker-compose.yaml이다. 우선 services를 train0, tensorboard0, train1, tensorboard1 이렇게 4개를 만들었다. train0는 GPU0을 사용하고, train1은 GPU1을 사용한다. tensorboard0는 포트 6006을 이용해 tensorboard를 실행하였고, tensorboard1은 포트 6005를 이용해 tensorboard를 실행한다. 이렇게 세팅하여 docker-compose를 실행하면 두 개의 코드가 충돌없이 아주 잘 돌아간다. 그리고 tensorboard가 참조하는 logdir도 다르게 설정해줘야한다. 저장할때부터 다른 경로로 따로 저장해야한다.

이렇게 동시에 두 모델을 학습하며 두 개의 tensorboard를 띄울 수 있다.

 

참고:

haanjack.github.io/docker/2018/04/14/nvidia-gpu-isolation.html

stackoverflow.com/questions/43368470/use-nvidia-docker-from-docker-compose