JAX를 사용하여 Cloud TPU VM에서 계산 실행

gcloud를 사용하여 Cloud TPU VM 만들기

1. 명령어를 더 쉽게 사용할 수 있도록 몇 가지 환경 변수를 정의합니다.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   환경 변수 설명

   변수	설명
   PROJECT_ID	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
   TPU_NAME	TPU의 이름입니다.
   ZONE	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
   ACCELERATOR_TYPE	가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
   RUNTIME_VERSION	Cloud TPU 소프트웨어 버전입니다.

2. Cloud Shell 또는 Google Cloud CLI가 설치된 컴퓨터 터미널에서 다음 명령어를 실행하여 TPU VM을 만듭니다.

   $ gcloud compute tpus tpu-vm create $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --accelerator-type=$ACCELERATOR_TYPE \
       --version=$RUNTIME_VERSION

Cloud TPU VM에 연결

다음 명령어를 사용하여 SSH를 통해 TPU VM에 연결합니다.

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

SSH를 사용하여 TPU VM에 연결할 수 없는 경우 TPU VM에 외부 IP 주소가 없기 때문일 수 있습니다. 외부 IP 주소가 없는 TPU VM에 액세스하려면 공개 IP 주소가 없는 TPU VM에 연결의 안내를 따르세요.

Cloud TPU VM에 JAX 설치

(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

시스템 확인

JAX가 TPU에 액세스할 수 있고 기본 작업을 실행할 수 있는지 확인합니다.

1. Python 3 인터프리터 시작:

   (vm)$ python3

   >>> import jax

2. 사용 가능한 TPU 코어 수 표시:

   >>> jax.device_count()

TPU 코어 수가 표시됩니다. 표시되는 코어 수는 사용 중인 TPU 버전에 따라 다릅니다. 자세한 내용은 TPU 버전을 참조하세요.

계산 수행

>>> jax.numpy.add(1, 1)

numpy add의 결과가 표시됩니다.

명령어에서 출력합니다.

Array(2, dtype=int32, weak_type=True)

Python 인터프리터 종료

>>> exit()

TPU VM에서 JAX 코드 실행

이제 원하는 JAX 코드를 실행할 수 있습니다. Flax 예시는 JAX에서 표준 ML 모드 실행을 시작할 수 있는 훌륭한 장소입니다. 예를 들어 기본 MNIST 컨볼루션 네트워크를 학습시키려면 다음 안내를 따르세요.

1. Flax 예시 종속 항목 설치

   (vm)$ pip install --upgrade clu
   (vm)$ pip install tensorflow
   (vm)$ pip install tensorflow_datasets

2. Flax를 설치합니다.

   (vm)$ git clone https://github.com/google/flax.git
   (vm)$ pip install --user flax

3. Flax MNIST 학습 스크립트를 실행합니다.

   (vm)$ cd flax/examples/mnist
   (vm)$ python3 main.py --workdir=/tmp/mnist \
       --config=configs/default.py \
       --config.learning_rate=0.05 \
       --config.num_epochs=5

스크립트가 데이터 세트를 다운로드하고 학습을 시작합니다. 스크립트 출력은 다음과 같아야 합니다.

I0214 18:00:50.660087 140369022753856 train.py:146] epoch:  1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88
I0214 18:00:52.015867 140369022753856 train.py:146] epoch:  2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72
I0214 18:00:53.377511 140369022753856 train.py:146] epoch:  3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04
I0214 18:00:54.727168 140369022753856 train.py:146] epoch:  4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15
I0214 18:00:56.082807 140369022753856 train.py:146] epoch:  5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18