[Pytorch] Multi GPU를 활용 해 보자

일반적으로 하나의 gpu를 사용할 경우 아래와 같이 이용 가능

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = Net()
net.to(device)

multi-gpu를 이용하기 위해서는 아래와 같이 nn.DataParallel을 이용하여 손쉽게 사용 가능하다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = Net()
if torch.cuda.device_count() > 1:
    net = nn.DataParallel(net)
net.to(device)

이런식으로 설정 할 경우 선언한 batch size가 알아서 각 gpu로 분배되는 방식.

여러 gpu를 효율적으로 사용하려면 사용할 gpu 수만큼 batch size도 늘려주자.

추가적으로 어느 gpu를 사용할 것인가에 대한 선택의 방법은 확인 후 수정. 아래의 방법을 사용.

pytorch 기능을 써서 설정하는 방법은 모르겠다. 

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