torch - GPU 사용하기

일반적으로 하나의 GPU를 사용할 때는 다음과 같은 코드를 사용한다.

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = Net()
model.to(device)

 

하지만 다수의 GPU를 사용한다면 다음 코드와 같이 nn.DataParallel을 이용한다.

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

 nn.DataParallel을 사용할 경우 배치 크기가 알아서 각 GPU로 분배되는 방식으로 작동하기 때문에 GPU 수 만큼 배치 크기도 늘려 주어야 한다.