[딥러닝] 파이토치 multi GPU 로 학습시키기 (feat. Pytorch lightning )

딥러닝에서 여러개의 gpu를 이용해서 학습을 시키면 학습에 걸리는 시간을 줄일 수 있습니다. Pytorch를 이용해서 코딩을 해도 되고 Pytorch lightning을 이용해서 보다 간편하게 학습을 시킬 수도 있습니다. 본 글에서는 파이토치 라이트닝을 이용해서 Multi-GPU로 학습시킬 때 주의할 점에 대해서 말씀드리겠습니다. 병렬학습이라고도 할 수 있습니다.

 

확인사항

우선 시스템에서 GPU를 잘 인지하는지 확인해봅시다. 만약 에러가 날 경우에는 많은 경우 올바른 버전의 패키지를 설치하지 않았기 때문입니다. 우선 쥬피터 노트북이든 그냥 파이썬이든 열어서 아래와 같은 명령어를 입력해봅시다.

 

import torch, os

#보통 이 두 명령어면 됩니다.

#GPU 인지 되는지 확인
print(torch.cuda.is_available())
#입력값이 True이면 cuda 인지, 즉 GPU를 이용한 학습 가능
# 몇 개 인지 출력
print(torch.cuda.device_count())
# 아래 명령어도 이용할 수 있습니다.
# 시스템에 있는 CUDA를 출력합니다.
print(os.environ['CUDA_VISIBLE_DEVICES'])

 

만약 gpu 있는거 확실한데, 인지가 안된다면 NVIDIA-SMI ,nvcc--version을 통해서 현재 시스템의 cuda버전을 확인합니다.  여기에 보이는 cuda 버전 과 conda list | grep torch 를 해서 현재 활성화된 환경에서 설치된 파이토치가 cuda 버전 몇 번째인지 확인합니다. 필요에 따라서 다시 설치해줍니다. 

 

기존 코드를 multi GPU학습에 맞게 고치기

이제 기존 single GPU 나 CPU 기준이었던 학습코드가 있다면 multi-GPU를 위해서 코드를 바꾸어줍니다. 이 블로그는 파이토치 라이트닝을 기준으로 공식 홈페이지에 나온 내용을 알려드리고 있습니다. 만약 파이토치 라이트닝 아니고 그냥 파이토치를 하고 있다면 여기서는 안 다룹니다. 파이토치 라이트닝이 조금 짜증날 때도 있지만 여러 기능들이 편한 건 사실이라서요. 

 

1. 기존 코드에 있던 .to('cuda') 또는 .to(device)를 모두 지웁니다. 
   • model.to(device), X_train.to(device),y_train.to(device) 등등이요.
2. 원래 data module 에 train_dataloader를 지정했거나 그냥 Trainer 에 train_dataloader를 입력값으로 주었다면 Lightning 모듈 안으로 넣어주세요. train_dataloader, val_dataloader, test_dataloader 모두요.  아래 코드 예시를 확인해주세요.
3. Trainer 에서 devices에 gpu 갯수 (torch.cuda.device_count()) 를 적어줍니다.

def train_dataloader(self):
	return DataLoader(self.train_data, batch_size=self.batch_size, drop_last=True, shuffle=True)
    
def test_dataloader(self):
	return DataLoader(self.test_data, batch_size=self.batch_size, shuffle=False)

 

의외로 간단하죠?

이렇게 병렬학습과 그 전글에서 언급한 mixed_precision training을 합치면 학습시간을 줄일 수 있습니다. 

파이토치 라이트(Pytorch lightning)에서 이용하기 좋은 팁

저는 텐서보드 로거를 이용하는 걸 좋아합니다. 또한 처음에는 profiler를 이용해서 각 단계에 들어가는 시간을 확인하는 것도 좋아합니다.

• 파이토치 라이트닝에서 모델 모양을 그래프로 보고 싶으시면 log_graph=True 를 해주시면 됩니다. 보다 정확하게 말씀드리면 Trainer(logger=TensorBoardLogger(log_graph=True)입니다). 또한 클래스에서 self.example_input_array를 torch.rand(1, 데이터 모양) 으로 지정해주면 됩니다. 초반에 조금 시간이 걸려도 한번에 보기 쉬워서 전 좋아합니다. 
• Trainer(profiler="simple") 을 추가하면 각 단계별로 걸리는 시간에 대해 알 수 있습니다. 초반에 쓰기 좋아요.
• 진행바는 보통 기본값으로 들어가 있는데 아니면 그냥 TQDMProgressBar를 callbacks에 추가해주세요.
• 콜백은 옵션입니다. 정말 다양한 콜백들이 있습니다. callbacks =[early_Stopping, checkpoint_callback, TQDMProgressBar(refresh_rate = 20)] 을 쓰는 것 같습니다. 스케쥴러나 파인튜닝 관련 콜백도 있습니다.

 

기타 딥러닝에서 이용하기 좋은 팁

몇 가지 뻘짓을 했던 기억이 있으니 몇 자 적어보겠습니다.

• 만약 지금 인풋 행렬의 열을 바꾸어야 하면 reshape보다는 transpose나 permute() 이용하세요. 예를 들어서 Convolution layer을 기준으로 파이토치(Pytorch)와 케라스(keras)는 hidden layer 가 오는 위치가 다릅니다. 파이토치는 [데이터 갯수, hidden channel, 데이터 길이] 이 순서입니다. 만약 인풋이 [데이터 갯수, 데이터 길이, hidden channel] 순서라면 reshape 보다는torch.transpose(X, 1,2) 를 이용해주세요. 1번째 차원과 2번째 차원을 transpose한다는 뜻입니다. 
• train_dataloader 에는 shuffle=True가 학습에 도움이 되는데 val_dataloader 나 test_dataloader에는 안 해주셔도 됩니다. shuffle = False 로 놔주세요
• 만약 loss를 보는데 중간에 스파이크가 보이면 해당 batch의 데이터 갯수가 너무 적어서 그럴 수 있으니 drop_last=True가 도움이 될 수 있습니다.
• 만약 classification을 위해서 crossEntropyLoss를 이용하신다면, label의 datatype 이 int인 걸 확인해주세요. 확인 방법은 y_label.dtype 이고 아니면 y_lable.type(torch.int64)으로 바꾸어주세요. 아 shape이 (size,1)인것도 확인해주세요.