Note ≡ | GPU 사용
일반적으로 하나의 GPU를 사용할 때는 다음과 같은 코드를 이용합니다.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = Net() model.to(device)
하지만 사용하는 PC에서 다수의 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을 사용할 경우 배치 크기(batch size)가 알아서 각 GPU로 분배되는 방식으로 작동합니다. 따라서 GPU 수만큼 배치 크기도 늘려 주어야 합니다.