컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.

Docker 컨테이너에서 Cloud TPU 애플리케이션 실행

Docker 컨테이너를 사용하면 코드와 필요한 모든 종속 항목을 배포 가능한 패키지 하나로 결합하여 애플리케이션을 쉽게 구성할 수 있습니다. TPU VM 내에서 Docker 컨테이너를 실행하면 Cloud TPU 애플리케이션을 간단하게 구성하고 공유할 수 있습니다. 이 문서에서는 Cloud TPU에서 지원하는 각 ML 프레임워크에서 Docker 컨테이너를 설정하는 방법을 설명합니다.

Docker 컨테이너에서 TensorFlow 모델 학습

TPU 기기

  1. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

    FROM python:3.8
    RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
    RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so
    RUN git clone https://github.com/tensorflow/models.git
    WORKDIR ./models
    RUN pip install -r official/requirements.txt
    ENV PYTHONPATH=/models
    
  2. Cloud Storage 버킷 만들기

    gsutil mb -c standard -l europe-west4 gs://your-bucket-name
    
  3. TPU VM을 만듭니다.

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-tf-2.11.0
    
  4. Dockerfile을 TPU VM에 복사합니다.

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
    
  5. TPU VM에 SSH를 통해 연결

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  6. Docker 이미지를 빌드합니다.

    sudo docker build -t your-image-name .
    
  7. Docker 컨테이너를 시작합니다.

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
    
  8. 환경 변수 설정하기

    export STORAGE_BUCKET=gs://your-bucket-name
    export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
    export ZONE=europe-west4-a
    export TPU_NAME=local
    
  9. ResNet을 학습시킵니다.

    python3 official/vision/train.py \
    --tpu=local \
    --experiment=resnet_imagenet \
    --mode=train_and_eval \
    --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
    --model_dir=${MODEL_DIR} \
    --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
    

TPU 포드

  1. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

    FROM python:3.8
    RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
    RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so
    RUN git clone https://github.com/tensorflow/models.git
    WORKDIR ./models
    RUN pip install -r official/requirements.txt
    ENV PYTHONPATH=/models
    
  2. TPU VM을 만듭니다.

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v3-32 \
    --version=tpu-vm-tf-2.11.0-pod
    
  3. Dockerfile을 TPU VM에 복사합니다.

    gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
    
  4. TPU VM에 SSH를 통해 연결

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  5. Docker 이미지를 빌드합니다.

    sudo docker build -t your-image-name .
    
  6. Docker 컨테이너를 시작합니다.

    sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
    
  7. ResNet을 학습시킵니다.

    python3 official/vision/train.py \
    --tpu=your-tpu-name \
    --experiment=resnet_imagenet \
    --mode=train_and_eval \
    --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
    --model_dir=${MODEL_DIR} \
    --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
    

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

  1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
  2. TPU VM을 종료하려면 exit를 입력합니다.
  3. TPU VM을 삭제합니다.
      $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

Docker 컨테이너에서 PyTorch 모델 학습

  1. Cloud TPU VM을 만듭니다.

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-pt-1.13
    
  2. TPU VM에 SSH를 통해 연결합니다.

    gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
    
  3. TPU VM에서 컨테이너를 시작합니다.

    sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
    
  4. libtpu를 설치합니다.

    pip install torch_xla[tpuvm]
    export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    
  5. PyTorch XLA 저장소를 클론합니다.

    git clone --recursive https://github.com/pytorch/xla.git
    
  6. ImageNet을 학습시킵니다.

    python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
    

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

  1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
  2. TPU VM을 종료하려면 exit를 입력합니다.
  3. TPU VM을 삭제합니다.
    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

Docker 컨테이너에서 JAX 모델 학습

TPU 기기

  1. TPU VM 만들기

    gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=europe-west4-a \
    --accelerator-type=v2-8 \
    --version=tpu-vm-base
    
  2. TPU VM에 SSH를 통해 연결합니다.

    gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
    
  3. TPU VM에서 Docker 데몬을 시작합니다.

    sudo systemctl start docker
    
  4. Docker 컨테이너 시작

    sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
    
  5. JAX를 설치합니다.

    pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    
  6. FLAX를 설치합니다.

    pip install --upgrade clu
    git clone https://github.com/google/flax.git
    pip install --user -e flax
    
  7. FLAX MNIST 학습 스크립트 실행

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

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

  1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
  2. TPU VM을 종료하려면 exit를 입력합니다.
  3. TPU VM을 삭제합니다.

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
    

TPU 포드

TPU Pod에서 JAX 코드를 실행하는 경우 모든 TPU 워커에서 JAX 코드를 동시에 실행해야 합니다. 이를 위한 한 가지 방법은 gcloud compute tpus tpu-vm ssh 명령어를 --worker=all--command 플래그와 함께 사용하는 것입니다. 다음 절차는 각 TPU 워커를 더 쉽게 설정할 수 있도록 Docker 이미지를 만드는 방법을 보여줍니다.

  1. 터미널 프롬프트에서 다음 환경 변수를 설정합니다.

    export USER=your-user-id
    export PROJECT=your-project-name
    export REPO_NAME=your-repo-name
    export IMAGE_NAME=your-image-name
    export TAG=your-tag
    export TPU_NAME=your-tpu-name
    export ZONE=europe-west4-a
    export ACCEL_TYPE=v2-8
    

    다음과 같이 바꿉니다.

    • your-user-id를 사용자 ID로 바꿉니다.
    • your-project-name을 Google Cloud 프로젝트 이름으로 바꿉니다.
    • your-repo-nameArtifact Registry 저장소 이름으로 바꿉니다.
    • your-image-name을 Docker 이미지 이름으로 바꿉니다.
    • your-tag를 Docker 이미지의 태그로 바꿉니다.
    • your-tpu-name을 TPU 이름으로 바꿉니다.
  2. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

    FROM google-cloud-cli:latest
    RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    RUN pip install --upgrade clu
    RUN git clone https://github.com/google/flax.git
    RUN pip install --user -e flax
    WORKDIR ./flax/examples/mnist
    
  3. Docker 이미지를 빌드합니다.

    docker build -t $IMAGE .
    
  4. Artifact Registry로 푸시하기 전에 Docker 이미지에 태그를 추가합니다. Artifact Registry 작업에 대한 자세한 내용은 컨테이너 이미지로 작업하기를 참조하세요.

    docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
    
  5. Docker 이미지를 Artifact Registry로 푸시합니다.

    docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
    
  6. TPU VM을 만듭니다.

    gcloud compute tpus tpu-vm create $TPU_NAME \
    --zone $ZONE \
    --accelerator-type $ACCEL_TYPE \
    --version tpu-vm-base
    
  7. 모든 TPU 워커에서 Artifact Registry의 Docker 이미지를 가져옵니다.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="sudo usermod -a -G docker $USER"
    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="docker pull europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag"
    
  8. 모든 TPU 워커에서 컨테이너를 실행합니다.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag bash"
    
  9. 모든 TPU 워커에서 학습 스크립트를 실행합니다.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
    --zone $ZONE \
    --command="docker exec --privileged $CONTAINER_NAME python3 main.py --workdir=/tmp/mnist \
    --config=configs/default.py \
    --config.learning_rate=0.05 \
    --config.num_epochs=5"
    

학습 스크립트가 완료되면 리소스를 삭제해야 합니다. 다음 명령어를 사용하여 TPU VM을 삭제합니다.

  $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
  

다음 단계