[PyTorch] 특정 GPU로 정해 동작시키기

<분류>
A. 수단
- OS/Platform/Tool : Linux, Kubernetes(k8s), Docker, AWS
- Package Manager : node.js, yarn, brew, 
- Compiler/Transpillar : React, Nvcc, gcc/g++, Babel, Flutter
- Module Bundler  : React, Webpack, Parcel

B. 언어
- C/C++, python, Javacsript, Typescript, Go-Lang, CUDA, Dart, HTML/CSS

C. 라이브러리 및 프레임워크 및 SDK
- OpenCV, OpenCL, FastAPI, PyTorch, Tensorflow, Nsight

 

---

1. What? (현상)

PyTorch를 활용할 때, GPU 중에서 어떤 GPU를 특정해서 돌리고 싶을 때 어떻게 지정해서 돌리면 좋을지 살펴보려고합니다. 

 

한 개의 GPU 일 때와 여러개의 GPU일 때를 나누어 살펴보도록 합니다.

---

2. Why? (원인)

• X

---

3. How? (해결책)

 

먼저 가능한 GPU를 아래와 같은 명령어로 확인합니다.

import torch

print([torch.cuda.device(i) for i in range(torch.cuda.device_count())])

[<torch.cuda.device object at 0x7f0e5680e610>,
<torch.cuda.device object at 0x7f0e5680e611>,
<torch.cuda.device object at 0x7f0e5680e612>,
<torch.cuda.device object at 0x7f0e5680e613>,
<torch.cuda.device object at 0x7f0e5680e614>,
<torch.cuda.device object at 0x7f0e5680e615>,
<torch.cuda.device object at 0x7f0e5680e616>,
<torch.cuda.device object at 0x7f0e5680e617>]

• 하나의 GPU로 돌리는 방법
  • 0~7번의 GPU들 8개 중에 7번 GPU를 돌리고 싶은 상황이라고 가정.

class encoder(torch.nn.Module):
    def __init__(self, ...):
    	...
        self = self.cuda
    def forward(self, ...) :

torch.cuda.set_device(7)
print(torch.cuda.current_device())

• 여러개의 GPU로 돌리는 방법1. Data Parallel에서 명시해주기
  • 0~7번의 GPU들 8개 중에 [3,4,7] GPU를 돌리고 싶은 상황이라고 가정.
  • 단, 이 방법은 0번 GPU가 동작하지 않는 경우 정확히 3,4,7번으로 동작하지 않을 수 있습니다. 이런 경우, 아래의 방법을 선택해주세요.

cuda = torch. cuda.is_available()
device = torch.device(‘cuda’ if cuda else ‘cpu)

model = torch.nn.DataParallel(model, device_ids=[3,4,7]).to(device)

• 여러개의 GPU로 돌리는 방법2. 실행시 명시해주기
  • 0~7번의 GPU들 8개 중에 [3,4,7] GPU를 돌리고 싶은 상황이라고 가정.
  • 위와 다르게 3,4,7로 명시하는 것이 아닌 0,1,2로 명시해주어야 합니다.

CUDA_VISIBLE_DEVICES=3,4,7 python3 main.py

import os

cuda = torch. cuda.is_available()
device = torch.device(‘cuda’ if cuda else ‘cpu)

model = torch.nn.DataParallel(model, device_ids=[0,1,2]).to(device)

print(os.environ.get(‘CUDA_VISIBLE_DEVICES’))

---