[Pytorch] OS에 따른 GPU 자동설정 코드

Pytorch 코드를 이리저리 다른 환경에서 실행해야할 일이 많다.
Linux 서버에서 실행해야할 때도 있고, Mac 내에 내장 GPU로 실행할 때, Window로 간단히 print만 찍어볼 때, 등등 다양한 환경에서 작업을 했다.

하지만 코드를 실행할 때마다 오류가 발생했는데 오류가 발생하는 곳은 Pytorch에서 GPU 사용을 위해 정의하는 변수인 device 변수에서 발생했다.

Window나 Linux에서는 cuda를 사용하지만 Mac에서는 mps를 사용하고, 이를 체크하는 메소드 또한 다르기 때문에 다른 환경에서 실행할 떄 계속 변경해줘야 했다.

이것을 User PC의 OS를 자동으로 인식해서 돌아가게 만들면 어떨가해서 코드를 찾아보았다.

import platform

# OS Check
os_name = platform.system()
if os_name == 'Windows':
    OS= 'Windows'
    print(os_name)
elif os_name == 'Darwin':
    OS= 'Mac'
    print(os_name)
elif os_name == 'Linux':
    OS= 'Linux'
    print(os_name)
else:
    print(f'알 수 없는 운영 체제: {os_name}')

Mac OS 환경에서 해당 코드를 실행한 결과이다.
platform라이브러리를 통해서 코드가 실행되는 User PC의 OS를 알 수 있다.

# 윈도우와 리눅스이면 torch에서 cuda를 사용할 수 있도록 설정
if OS == 'Windows' or OS == 'Linux':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'운영체제: {OS}, 디바이스: {device}')
elif OS == 'Mac':
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") # Mac
    print(f'운영체제: {OS}, 디바이스: {device}')
else:
    print(f'알 수 없는 운영 체제: {OS}')

윈도우나 리눅스에는 그래픽카드를 사용하기 때문에 torch.cuda.is_available() 메소드를 사용하여 사용 가능 여부를 확인했고, 맥에서는 torch.backends.mps.is_available()를 통해 내장 GPU 사용 가능여부를 확인했다.

맥북프로 환경에서 코드를 실행한 결과는 다음과 같았다.

하나의 OS만 사용한다면 굳이 이런 코드를 사용하지 않아도 되겠지만, 팀 협업에서 내가 작성한 코드 실행에서 조금이라도 오류를 줄이기 위해 이런 템플릿 하나 있어도 괜찮지 않을까 생각하여 이 글을 작성했습니다.
감사합니다.