특정 GPU만 사용하기

GPU MIG 모드를 사용했을때 특정 GPU만 사용하고 싶을 때는 일단 먼저 MIG UUID가 필요하다

$ nivida-smi -L

  MIG 2g.20gb     Device  0: (UUID: MIG-)
  MIG 2g.20gb     Device  1: (UUID: MIG-)
  MIG 2g.20gb     Device  2: (UUID: MIG-)

Python

CUDA_VISIBLE_DEVICES="<MIG UUID>" python train.py

Jupyter Notebook

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" 
os.environ["CUDA_VISIBLE_DEVICES"]= "<MIG-UUID>"

 

torch를 import 하기 전에 먼저 환경을 설정 해야된다

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Docker Compose

  gpu:
    environment:
      CUDA_DEVICE_ORDER: "PCI_BUS_ID"
      NVIDIA_VISIBLE_DEVICES: "<MIG UUID>"
      NVIDIA_DRIVER_CAPABILITIES: compute,utility
    runtime: nvidia 
    restart: always