[JAX] JAX 최신버전(0.8.0~0.9.0) GPU 체크

 

본인의 컴퓨터나 노트북에 NVIDIA GPU가 있고 CUDA가 이미 설치되어있는 것을 전제로 한다.

또한 가상환경을 conda에서 만들어서 그 안에서 최신 버전의 JAX를 설치하는 과정이다. (일반적으로 JAX를 pip으로 설치하면 무조건 최신버전을 설치하도록 된다.)

 

설치한 JAX 버전은 현재 기준으로 가장 최신인 0.9.0이다. 이전 버전을 설치하고 싶다면 다음 글이 더 유용하다.

 

 

나는 주로 Linux x86_64, Windows WSL2 환경을 사용한다.

 

 

다음 커맨드를 실행하면 JAX GPU 버전이 설치가 되는데 여기서 가정하는 것은 CUDA가 적절하게 설치되어있다는 것이다.

pip install --upgrade "jax[cuda12]"

 

 

최근에는 JAX installation 문서가 잘 업데이트되어서 설치하는 데에 크게 어려움은 없다. 다만 설치하고 gpu를 잘 사용되는지 확인하고자 한다면 다음 처럼 확인해보자.

 

python # ubuntu 환경에서 python 실행
>> import jax
>> print( jax.default_backend() )

 

 

그러면 CPU를 사용한다고 나온다. 이런 문제는 보통 CUDA가 전체 컴퓨터에 설치되어있다하더라도 내가 셋업하고 있는 환경 안에서 CUDA를 읽어들이지 못하는 경우이다. 따라서 환경의 안정성을 위해서는 각 가상환경마다 NVIDIA CUDA를 설치해주는 게 좋긴 하다. (다만 이런 환경이 많이 늘어난 상태라 용량이 문제가 된다면 언젠가는 정리해야할 것 같기도...) 다음 커맨드가 가상환경 내에서 CUDA를 사용하기 위한 library까지 포함해서 설치하는 것이다.

 

(Note : 옛날에 윈도우 자체에서 CUDA 버전을 여러 개 설치해서 사용을 했었는데 이러면 각 CUDA와 가상환경 사이에 읽어들이는 게 꼬이기 쉬웠었다.)

 

CUDA를 사용하는지 확실하게 하기 위해서는 다음 Command도 다시 설치해준다.

 

이렇게 하면 cudnn을 비롯한 CUDA를 사용하기 위한 많은 라이브러리들을 설치하도록 되어있다.

 

 

 

 

이제는 GPU를 인식하는 것을 확인할 수 있다.