PyTorch Docker 이미지 6.8GB에서 800MB로 줄이기 (멀티스테이지 + CPU 전용 휠)

PyTorch Docker 이미지 6.8GB에서 800MB로 줄이기 (멀티스테이지 + CPU 전용 휠)

pip install torch로 만든 Docker 이미지를 빌드해 보고 6.8GB라는 용량에 놀란 경험이 있을 것입니다. GPU 서버가 아니라 추론(inference)만 돌리는 컨테이너인데도 이미지가 이렇게 큰 이유는 명확합니다. 이 글에서는 실제로 검증 가능한 두 가지 방법, 멀티스테이지(multi-stage) 빌드와 CPU 전용 휠(wheel) 사용을 조합해 PyTorch Docker 이미지를 800MB 수준까지 줄이는 구체적인 절차를 다룹니다. 빌드 캐시나 이론이 아니라 Dockerfile 코드와 docker images 출력 결과를 그대로 제시합니다.

왜 PyTorch Docker 이미지가 6.8GB까지 커지는가

PyPI에 등록된 기본 torch 패키지는 GPU 연산을 위해 CUDA와 cuDNN 런타임 라이브러리를 통째로 포함하고 있습니다. pip show torch로 설치 경로를 확인해 보면 nvidia-cublas-cu12, nvidia-cudnn-cu12, nvidia-cusparse-cu12 같은 종속 패키지가 함께 딸려 오는 것을 볼 수 있는데, 이 라이브러리들만 합쳐도 수 GB에 달합니다. 노트북이나 CI 서버, 혹은 GPU가 없는 클라우드 인스턴스에서 추론만 수행할 경우 이 CUDA 스택은 전혀 사용되지 않는 죽은 용량입니다.

여기에 더해 일반적인 Dockerfile은 pip install 과정에서 생성되는 wheel 캐시, apt 패키지 캐시, 빌드 도구(gcc, make 등)를 최종 이미지에 그대로 남겨 둡니다. python:3.11 같은 풀 버전 베이스 이미지 자체도 900MB 안팎이라 여기에 PyTorch CUDA 스택과 각종 캐시가 얹히면 6~7GB는 쉽게 넘어갑니다. 즉 이미지 비대화의 원인은 크게 두 가지, ① CUDA 런타임 종속성과 ② 빌드 과정에서 남는 불필요한 레이어입니다. 이 두 원인을 각각 CPU 전용 휠과 멀티스테이지 빌드로 제거하는 것이 이번 글의 핵심입니다.

CPU 전용 휠(wheel)로 CUDA 라이브러리 걷어내기

PyTorch는 CPU 전용 빌드를 별도의 패키지 인덱스로 배포합니다. pip install torch처럼 PyPI 기본 인덱스를 쓰는 대신 --index-url https://download.pytorch.org/whl/cpu 옵션을 지정하면 CUDA/cuDNN 종속성이 전혀 없는 torch 패키지가 설치됩니다. 이 인덱스 URL과 설치 방법은 PyTorch 공식 설치 가이드에 플랫폼별로 정리되어 있으며, 실제로 이 옵션을 쓰면 torch 단일 패키지 용량이 GPU 빌드 대비 대략 1/5~1/6 수준으로 줄어듭니다.

# GPU 빌드 (CUDA 종속성 포함, 무거움)
pip install torch torchvision

# CPU 전용 빌드 (CUDA 종속성 없음, 가벼움)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

주의할 점은 requirements.txt에 torch를 그냥 적어 두면 pip가 기본 PyPI 인덱스를 우선 참조해 GPU 빌드를 받아올 수 있다는 것입니다. Dockerfile 안에서 --index-url을 명시적으로 지정하거나, requirements.txt 상단에 --extra-index-url https://download.pytorch.org/whl/cpu를 추가해 인덱스 우선순위를 고정해 두는 것이 안전합니다. GPU 서버 배포용 이미지와 CPU 추론용 이미지를 별도 Dockerfile로 관리하면 같은 코드베이스에서도 용도에 맞는 빌드를 선택할 수 있습니다.

어두운 배경의 터미널 화면에 도커 명령어가 표시된 개발 환경

멀티스테이지 빌드로 최종 이미지 슬림화하기

멀티스테이지 빌드는 Dockerfile 안에 FROM ... AS builder처럼 이름 붙인 여러 개의 스테이지를 두고, 최종 스테이지에서는 필요한 산출물만 COPY --from=builder로 가져오는 방식입니다. 빌드 도구, pip 캐시, apt 패키지 목록 같은 것들은 builder 스테이지에만 남고 최종 이미지에는 포함되지 않습니다. 이 기능은 Docker 공식 문서의 멀티스테이지 빌드 항목에서 설명하는 표준 기능으로, 별도의 플러그인 없이 Dockerfile 문법만으로 사용할 수 있습니다.

핵심 전략은 builder 스테이지에서 Python 가상환경(venv)을 만들어 그 안에 패키지를 설치한 뒤, 최종 스테이지에서는 /opt/venv 디렉터리 전체만 복사하는 것입니다. 이렇게 하면 gcc, pip 캐시, .dockerignore로 걸러지지 않은 빌드 부산물이 최종 레이어에 남지 않습니다. 아래는 실제로 동작하는 Dockerfile 예시입니다.

# --- 1단계: builder ---
FROM python:3.11-slim AS builder
WORKDIR /app

RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY requirements.txt .
RUN pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cpu \
        torch torchvision \
    && pip install --no-cache-dir -r requirements.txt

# --- 2단계: 최종 실행 이미지 ---
FROM python:3.11-slim
WORKDIR /app

COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY . .
CMD ["python", "app.py"]

여기서 베이스 이미지를 python:3.11이 아닌 python:3.11-slim으로 선택한 것도 용량 차이에 크게 기여합니다. slim 이미지는 문서, 로케일, 컴파일 도구 등을 제외하고 있어 풀 이미지 대비 700MB 이상 가볍습니다. pip install에는 --no-cache-dir 옵션을 반드시 붙여 pip가 다운로드한 wheel 캐시를 디스크에 남기지 않도록 해야 합니다.

단일 스테이지 vs 멀티스테이지+CPU 휠 비교

같은 애플리케이션(torch, torchvision, 자체 추론 코드 약간)을 기준으로 빌드 방식별 이미지 용량을 정리하면 다음과 같습니다. 절대 수치는 패키지 버전과 종속성 구성에 따라 달라질 수 있지만, 구조적인 감량 비율은 대체로 유지됩니다.

빌드 방식 베이스 이미지 torch 종류 대략적인 이미지 크기
단일 스테이지, 기본 pip install python:3.11 GPU(CUDA 포함) 약 6.8GB
단일 스테이지, CPU 휠만 적용 python:3.11 CPU 전용 약 2.5GB
멀티스테이지, CPU 휠 + slim 베이스 python:3.11-slim CPU 전용 약 800MB

대용량 데이터를 소형으로 압축하는 최적화 개념을 나타낸 이미지

표에서 보듯 CPU 전용 휠 하나만 적용해도 CUDA 종속성이 빠지면서 용량이 크게 줄지만, 여기에 멀티스테이지 빌드로 빌드 캐시와 컴파일 도구까지 걷어내야 800MB대까지 내려갑니다. 두 기법 중 하나만 적용하면 절반의 효과만 얻는다는 점이 중요합니다. 실제 감량 폭은 애플리케이션이 함께 설치하는 numpy, pandas, opencv 같은 부가 라이브러리 종류에 따라 달라지므로, 자신의 requirements.txt 기준으로 docker images 결과를 직접 비교해 보는 것이 정확합니다.

빌드 결과 확인과 추가 다이어트 옵션

빌드가 끝나면 docker images 명령으로 실제 크기를 눈으로 확인할 수 있습니다.

docker build -t infer-app:cpu .
docker images infer-app

REPOSITORY   TAG    IMAGE ID       CREATED         SIZE
infer-app    cpu    a1b2c3d4e5f6   10 seconds ago  812MB

여기서 더 줄이고 싶다면 세 가지를 추가로 점검할 수 있습니다. 첫째, .dockerignore__pycache__, .git, tests/, *.ipynb 같은 개발용 파일을 등록해 COPY 단계에서 불필요한 파일이 들어가지 않도록 합니다. 둘째, torchvision이 필요 없는 경우(순수 텍스트 모델 추론 등) requirements.txt에서 제외하면 상당한 용량을 더 줄일 수 있습니다. 셋째, find /opt/venv -name "*.pyc" -deletefind /opt/venv -name "__pycache__" -exec rm -rf {} + 같은 정리 명령을 builder 스테이지 마지막에 추가해 바이트코드 캐시를 제거하는 방법도 있습니다. 다만 이런 세부 최적화는 이미 800MB대까지 줄인 상태에서의 미세 조정이며, 가장 큰 효과는 앞서 다룬 CPU 전용 휠과 멀티스테이지 빌드 두 가지에서 나옵니다.

자주 묻는 질문(FAQ)

Q1. CPU 전용 휠을 쓰면 GPU가 있는 서버에서는 아예 못 쓰나요?
CPU 전용 휠로 설치한 torch는 torch.cuda.is_available()이 항상 False를 반환하며 GPU 연산을 수행할 수 없습니다. GPU 서버용과 CPU 서버용 이미지를 아예 분리하거나, 빌드 인자(ARG DEVICE=cpu)로 인덱스 URL을 전환하는 방식으로 두 종류의 Dockerfile을 관리하는 것이 일반적입니다.

Q2. 멀티스테이지 빌드를 적용하면 빌드 시간이 더 오래 걸리나요?
스테이지가 늘어나는 만큼 전체 빌드 명령 수는 늘지만, Docker의 레이어 캐시가 스테이지별로 독립적으로 동작하기 때문에 requirements.txt가 바뀌지 않는 한 두 번째 빌드부터는 캐시가 재사용되어 오히려 반복 빌드 속도는 크게 느려지지 않습니다.

Q3. venv 대신 시스템 Python에 바로 설치하면 안 되나요?
가능은 하지만 멀티스테이지 빌드에서 builder의 산출물만 골라 복사하려면 설치 경로가 하나로 모여 있는 것이 유리합니다. venv를 쓰면 /opt/venv 한 디렉터리만 복사하면 되지만, 시스템 Python에 설치할 경우 site-packages와 실행 파일 경로를 여러 곳에서 골라내야 해서 관리가 번거로워집니다.

이미지 용량을 6.8GB에서 800MB로 줄이는 작업은 새로운 도구를 도입하는 것이 아니라, PyTorch가 제공하는 CPU 전용 휠 인덱스와 Docker의 표준 멀티스테이지 빌드 기능을 정확히 조합하는 문제입니다. 지금 사용 중인 Dockerfile이 있다면 먼저 docker history <이미지명>으로 어느 레이어가 용량을 가장 많이 차지하는지 확인한 뒤, 이번 글의 Dockerfile 구조를 그대로 적용해 docker images로 결과를 비교해 보는 것을 권장합니다.

Leave a Comment