내가 필요해서 정리하는 JAX(2)(JAX의 설치 / 직접 설치해보고 적는 방법)

Docker

내가 가지고 있는 nvidia GPU의 하위버전이면 모든 환경세팅이 가능하다

내가 사용하는 GPU는 cuda 11.7 version의 nvidia gpu이다, 즉 cuda 11.7 이하의 cuda version의 도커는 모두 사용 가능하다.

 

내가 사용하는 GPU는 cuda 11.7 version의 nvidia gpu이다, 즉 cuda 11.7 이하의 cuda version의 도커는 모두 사용 가능하다.

 

 

 

해당 docker는 cuda 11.7을 기준으로 설정된 docker image이므로, 만약 하위 버전의 cuda를 사용한다면 google이나 docker hub를 통해 다운받으면 된다.

Dockerfile

0.00MB

Dockerfile을 내가 작업하고자 하는 디렉토리에 복사한 후, 다음 명령어로 image를 bulid한다.

docker build --tag jax_test .

 

이후, 아래와 같이 쉘 스크립트를 하나 만들어서 실행하면 된다. 스크립트의 파라미터들은 각자의 환경에 맡게 설정해준다.

docker run --rm -it \
	--gpus '"device=0"' \
	--name jax_cuda_11_7 \
	--workdir /workspace \
	-v /mnt:/mnt \
	-v container와 연결할 디렉토리:/workspace \
	--shm-size 16G \
	-p 8007:8007 \
	jax_test:latest /bin/bash

 

 

Conda

conda create -n jax python=3.9

conda activate jax

pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

차례대로 입력해주면 된다. (cuda 11 version)

다만, Conda로 환경세팅을 할 경우, 개인의 GPU cuda 버전에 따라 설정을 찾아서 해줘야 한다. 

 

 

 

++ 다음은 오늘 세팅한 환경으로 직접 모델을 돌려보려 한다.