파이토치에서 GPU를 사용할 수 있는지 어떻게 확인하나요?

공식 홈페이지StackOverflow 등에서 자주 보이는 질문과 답변을 번역하고 있습니다.

다음 링크에서 원문을 함께 찾아보실 수 있습니다.


질문

  • nvidia-smi 명령어로 사용 가능한 GPU가 있는 것은 확인하였습니다.
  • PyTorch에서는 사용 가능한 GPU를 어떻게 확인할 수 있나요?

답변

  • 다음과 같은 방법으로 GPU와 관련된 내용들을 알아볼 수 있습니다.
    >>> import torch # PyTorch를 불러옵니다.
    >>> torch.cuda.is_available() # GPU가 사용 가능한지 확인합니다.
    True
    >>> torch.cuda.device_count() # 사용 가능한 장치가 몇 개인지 확인합니다.
    2
    >>> torch.cuda.get_device_name(0) # 첫번째 GPU의 장치명을 확인합니다.
    'GeForce GTX TITAN X'
    >>> torch.cuda.get_device_name(1) # 두번째 GPU의 장치명을 확인합니다.
    'GeForce GTX TITAN X'
    
  • 다음과 같이 CPU와 GPU를 자동으로 선택하여 동작하도록 코드를 작성할 수 있습니다.
    import torch
    
    # GPU가 사용 가능하면 cuda, 아니면 cpu
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # .to() 메소드를 사용하여 텐서나 모델 등을 장치로 이동
    model = NeuralNetwork().to(device)
    

더 알아보기