본문 바로가기
AI

[PyTorch] M1, M2 칩으로 Pytorch GPU 사용하기

by jjjaeunn 2023. 8. 10.

글쓴이는 현재 M2 맥북 에어 환경에서 딥러닝 학습을 진행하고 있다. 본래 tensorflow 기반의 학습만 진행하다가 이번에 새로 pytorch 환경을 사용하게 되어서 Apple Silicon OS에서 pytorch 환경을 구축하게 되었다.

cuda 환경은 MacOS와 호환이 불가하기 때문에,  M1 M2 환경에서 pytorch를 사용할 수 있는 mps 장치 적용 방법에 대해서 알아보도록 하겠다.

 

💻 MPS란?

MPS 백엔드는 MacOS 장치에서 PyTorch 에코시스템을 확장하고 GPU에서 작업을 설정하고 실행할 수 있는 기존 스크립트 기능을 제공한다. 

1. torch version 확인하기

터미널을 키고, python과 torch 라이브러리 설치가 우선적으로 진행되어야 한다.

torch 1.12.0 이상의 버전부터 가능하므로, OS에 깔린 torch 버전을 확인한다.

 

#mps 빌드 확인
>>> torch.backends.mps.is_built()
#mps 사용 가능여부 확인
>>> torch.backends.mps.is_available()

PyTorch가 M2 칩을 사용하기 위한 MPS 장치 빌드가 되어있는 상태인지 + 사용이 가능한 상태인지를 확인한다.

각각 True를 출력하면 현재 mps 백엔드를 사용할 수 있는 상태인 것이다.

 

+) 참고로 이전 버전을 사용하면 mps를 못 찾는 오류가 발생한다. 

RuntimeError: Expected one of cpu, cuda, xpu, mkldnn, opengl, opencl, ideep, hip, ve, ort, mlc, xla, lazy, vulkan, meta, hpu device type at start of device string: mps

 

2. Torch 코드 device에 mps 적용하기

#mps device 사용하기
device = torch.device("mps")
print(f"current device is {device}")

쉘에서 이전에 device를 mps로 적용한 후, 이후의 데이터를 불러올 때에도 해당 device로 옮겨주는 작업이 필요하다.

이제 M2 로컬 환경에서도 gpu 자원을 사용할 수 있게 되었다 !

 


*참고

https://pytorch.org/docs/stable/notes/mps.html

https://discuss.pytorch.kr/t/mps/945/6