[Docker] Container 내부에서 GPU 사용하는 방법

안녕하세요, 오늘은 Docker Container 내부에서 GPU를 간단하게 사용할 수 있는 방법에 대해 소개하려고 합니다.

Docker Container에서 cuda가 실행되지 않는 모습

총 2가지를 진행하면 됩니다. 간단하게 nvidia-docker 설치와 같은 방법을 사용하지 않고 진행합니다.

1. Dockerfile 환경변수 추가 설정
2. Docker Container 실행시 매개변수 추가

 

 

Dockerfile 환경변수 추가 설정

---

# Dockerfile 

FROM python:3.8

ADD . /workdir

RUN apt-get update && apt-get install vim -y
RUN pip install torch && pip install transformers && pip install gpustat

##### 다음 줄 추가 #####
ENV CUDA_VISIBLE_DEVICES=0
##### 추가 #####

ENV CUDA_VISIBLE_DEVICES=0

 

다음과 같이 환경변수를 설정합니다.

 

Docker Container 실행시 매개변수 추가

---

위의 Dockerfile을 빌드합니다.(docker build -t <이미지 이름> <Dockerfile위치>)

 

# Container 생성
>>> docker run -dt --name test --gpus all image:latest

# Container 접속
>>> docker exec -it <container id> /bin/bash

여기서 컨테이너를 생성시 --gpus all 다음과 같은 매개변수를 추가합니다.

다음과 같이 GPU를 사용할 수 있게 됩니다. 

간단하게 로컬 GPU를 사용할 수 있도록 Docker image build 및 Container 생성에 대해 알아봤습니다.