[python Torch] 얼렁뚱땅 torch multigpu 쓰는 방법

안녕하세요 오랜만에 에러고치기로 돌아왔습니다 ! 

 

이번에는 에러라기보다는, MultiGPU를 쓰는 방법을 적어볼까 합니다

먼저 Ubuntu terminer 창에서 본인이 사용가능 한 GPU를 확인하는 방법은

 

 

이것으로 확인 가능합니다.

아래 터미널창을 보면 저는 GPU가 2개가 있으며, GPU ID는 0과 1입니다.

현재, 0과 1을 둘다 사용 중인 것도 확인 가능합니다.

nvidia-smi

 

만약 하나의 GPU만 사용할 것이라면 별도의 코드 변경 없이

터미널 창에서 어떤 GPU를 사용할 것인지 선택 가능합니다.

 

아래와 같이 터미널 창에 입력하면 됩니다.

여기서는 0번 GPU를 사용해서 main.py 파일을 돌리겠다라는 것입니다.

CUDA_VISIBLE_DEVICES=0 python main.py

 

 

그런데, MultiGPU를 사용하기 위해선 코드도 살짝 바꿔주어야 하는데요

한번 보겠습니다.

 

아래 코드는 저는 따로 Model.py라는 파일에 GeneratorUNet()이라는 모델을 정의해주었고

이를 generator라는 변수에 model을 정의해준 것입니다.

그러고, generator.cuda()를 통해 기존 cpu상에 있는 것을 GPU상에 올려줍니다.

generator = GeneratorUNet()
generator.cuda()

 

위를 변경하여 _generator에 임시로 모델 정의를 해주고,

그 다음 nn.DataParallel을 이용하여, MultiGPU가 가능하도록 해줍니다

요거를 안해주면 MultiGPU가 안되니 체크 반드시 하셔야 합니다 !

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]= "0,1"

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

_generator = GeneratorUNet().cuda()
generator = nn.DataParallel(_generator).to(device)

 

그럼 오늘은 여기서 끝 뿅