Menggunakan TPU untuk melatih model

Tensor Processing Unit (TPU) adalah ASIC Google yang dikembangkan secara khusus dan digunakan untuk mempercepat workload machine learning. Anda dapat menjalankan tugas pelatihan di Pelatihan AI Platform menggunakan Cloud TPU. AI Platform Training menyediakan antarmuka pengelolaan tugas sehingga Anda tidak perlu mengelola TPU sendiri. Sebagai gantinya, Anda dapat menggunakan AI Platform Training jobs API dengan cara yang sama seperti menggunakannya untuk pelatihan di CPU atau GPU.

TensorFlow API tingkat tinggi membantu Anda menjalankan model di hardware Cloud TPU.

Menyiapkan lingkungan Google Cloud Anda

Konfigurasikan lingkungan Google Cloud Anda dengan menyelesaikan bagian penyiapan pada panduan memulai.

Memberikan otorisasi pada Cloud TPU untuk mengakses project Anda

Ikuti langkah-langkah berikut untuk memberi otorisasi nama akun layanan Cloud TPU yang terkait dengan project Google Cloud Anda:

1. Dapatkan nama akun layanan Cloud TPU dengan memanggil projects.getConfig. Contoh:

   PROJECT_ID=PROJECT_ID
   
   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
       https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
   

2. Simpan nilai kolom serviceAccountProject dan tpuServiceAccount yang ditampilkan oleh API.

3. Lakukan inisialisasi akun layanan Cloud TPU:

   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
     -H "Content-Type: application/json" -d '{}'  \
     https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
   

Sekarang, tambahkan akun layanan Cloud TPU sebagai anggota di project Anda, dengan peran Cloud ML Service Agent. Selesaikan langkah-langkah berikut di Konsol Google Cloud atau menggunakan perintah gcloud:

Konsol

1. Login ke Google Cloud Console dan pilih project tempat Anda menggunakan TPU.
2. Pilih IAM & Admin > IAM.
3. Klik tombol Add untuk menambahkan anggota ke project.
4. Masukkan akun layanan TPU di kotak teks Anggota.
5. Klik daftar dropdown Peran.
6. Aktifkan peran Cloud ML Service Agent (Service Agents > Cloud ML Service Agent).

gcloud

1. Tetapkan variabel lingkungan yang berisi project ID Anda dan akun layanan Cloud TPU:

   PROJECT_ID=PROJECT_ID
   SVC_ACCOUNT=[e-mail]
   

2. Berikan peran ml.serviceAgent ke akun layanan Cloud TPU:

   gcloud projects add-iam-policy-binding $PROJECT_ID \
       --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
   

Untuk mengetahui detail selengkapnya tentang memberikan peran ke akun layanan, lihat dokumentasi IAM.

Contoh: Melatih model MNIST sampel

Bagian ini menunjukkan cara melatih contoh model MNIST menggunakan TPU dan runtime versi 2.11. Contoh tugas menggunakan tingkat skala BASIC_TPU yang telah ditetapkan untuk konfigurasi mesin Anda. Bagian selanjutnya dari panduan ini menunjukkan cara menyiapkan konfigurasi kustom.

Contoh ini mengasumsikan bahwa Anda menggunakan shell Bash dengan gcloud CLI terinstal. Jalankan perintah berikut untuk mendapatkan kode dan mengirimkan tugas pelatihan ke AI Platform Training:

1. Download kode untuk model referensi TensorFlow dan buka direktori yang berisi kode contoh:

   git clone https://github.com/tensorflow/models.git \
     --branch=v2.11.0 \
     --depth=1
   
   cd models
   

2. Buat file setup.py di direktori models. Ini akan memastikan bahwa perintah gcloud ai-platform jobs submit training menyertakan semua sub-paket yang diperlukan dalam direktori models/official saat membuat tarball kode pelatihan Anda, dan memastikan bahwa Pelatihan AI Platform menginstal Set Data TensorFlow sebagai dependensi saat menjalankan tugas pelatihan. Kode pelatihan ini mengandalkan set data TensorFlow untuk memuat data MNIST.

   Untuk membuat file setup.py, jalankan perintah berikut di shell Anda:

   cat << END > setup.py
   from setuptools import find_packages
   from setuptools import setup
   
   setup(
       name='official',
       install_requires=[
          'tensorflow-datasets~=3.1',
          'tensorflow-model-optimization>=0.4.1'
      ],
       packages=find_packages()
   )
   END
   

3. Kirim tugas pelatihan Anda menggunakan perintah gcloud ai-platform jobs submit training:

   gcloud ai-platform jobs submit training tpu_mnist_1 \
     --staging-bucket=gs://BUCKET_NAME \
     --package-path=official \
     --module-name=official.vision.image_classification.mnist_main \
     --runtime-version=2.11 \
     --python-version=3.7 \
     --scale-tier=BASIC_TPU \
     --region=us-central1 \
     -- \
     --distribution_strategy=tpu \
     --data_dir=gs://tfds-data/datasets \
     --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
   

   Ganti BUCKET_NAME dengan nama bucket Cloud Storage di project Google Cloud Anda. Gcloud CLI mengupload kode pelatihan terpaket ke bucket ini, dan AI Platform Training menyimpan output pelatihan di bucket.

4. Pantau tugas pelatihan Anda. Setelah tugas selesai, Anda dapat melihat outputnya di direktori gs://BUCKET_NAME/tpu_mnist_1_output.

Selengkapnya tentang melatih model di Cloud TPU

Bagian ini menjelaskan lebih lanjut cara mengonfigurasi tugas dan melatih model di Pelatihan AI Platform dengan Cloud TPU.

Menentukan region yang menawarkan TPU

Anda harus menjalankan tugas di wilayah tempat TPU tersedia. Region berikut saat ini menyediakan akses ke TPU:

• us-central1
• europe-west4

Untuk sepenuhnya memahami region yang tersedia untuk layanan Pelatihan AI Platform, termasuk pelatihan model dan prediksi online/batch, baca panduan wilayah.

Pembuatan versi Pelatihan TensorFlow dan AI Platform

Versi runtime AI Platform Training 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9, dan 2.11 tersedia untuk melatih model Anda di Cloud TPU. Lihat versi runtime Pelatihan AI Platform lebih lanjut dan versi TensorFlow yang sesuai.

Kebijakan pembuatan versi sama seperti untuk Cloud TPU. Dalam permintaan tugas pelatihan, pastikan untuk menentukan versi runtime yang tersedia untuk TPU dan cocok dengan versi TensorFlow yang digunakan dalam kode pelatihan Anda.

Menghubungkan ke server TPU gRPC

Dalam program TensorFlow, gunakan TPUClusterResolver untuk terhubung dengan server gRPC TPU yang berjalan di VM TPU.

Panduan TensorFlow untuk menggunakan TPU menunjukkan cara menggunakan TPUClusterResolver dengan strategi distribusi TPUStrategy.

Namun, Anda harus membuat satu perubahan penting saat menggunakan TPUClusterResolver untuk kode yang dijalankan di AI Platform Training: Jangan memberikan argumen apa pun saat Anda membuat instance TPUClusterResolver. Saat argumen kata kunci tpu, zone, dan project semuanya sudah ditetapkan ke nilai default None, AI Platform Training akan otomatis menyediakan cluster resolver dengan detail koneksi yang diperlukan melalui variabel lingkungan.

Contoh TensorFlow 2 berikut menunjukkan cara menginisialisasi cluster resolver dan strategi distribusi untuk pelatihan AI Platform Training:

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

Menggunakan TPU dalam kode TensorFlow

Untuk memanfaatkan TPU di komputer, gunakan TPUStrategy API TensorFlow 2. Panduan TensorFlow untuk menggunakan TPU menunjukkan cara melakukannya.

Untuk berlatih dengan TPU di TensorFlow 1, Anda dapat menggunakan TPUEstimator API. Panduan Cloud TPU untuk TPUEstimator API menunjukkan cara melakukannya.

Dokumentasi Cloud TPU juga menyediakan daftar operasi TensorFlow tingkat rendah yang tersedia di Cloud TPU.

Menggunakan TPU dalam kode PyTorch

Untuk memanfaatkan TPU saat Anda menggunakan container PyTorch yang telah dibangun sebelumnya, gunakan paket torch_xla. Pelajari cara menggunakan torch_xla untuk TPU dalam pelatihan di dokumentasi PyTorch. Untuk contoh penggunaan torch_xla lainnya, lihat tutorial di repositori GitHub PyTorch XLA

Perhatikan bahwa saat Anda berlatih menggunakan TPU di AI Platform Training, Anda menggunakan satu perangkat XLA, bukan beberapa perangkat XLA.

Lihat juga bagian berikut di halaman ini tentang cara mengonfigurasi tugas pelatihan untuk PyTorch dan TPU.

Mengonfigurasi mesin TPU kustom

Tugas pelatihan TPU berjalan pada konfigurasi dua VM. Satu VM (master) menjalankan kode Python Anda. Master ini menggerakkan server TensorFlow yang berjalan pada TPU worker.

Untuk menggunakan TPU dengan AI Platform Training, konfigurasikan tugas pelatihan Anda untuk mengakses mesin yang mendukung TPU melalui salah satu dari tiga cara:

• Gunakan tingkat skala BASIC_TPU. Anda dapat menggunakan metode ini untuk mengakses akselerator TPU v2.
• Gunakan pekerja cloud_tpu dan jenis mesin lama untuk VM master. Anda dapat menggunakan metode ini untuk mengakses akselerator TPU v2.
• Gunakan pekerja cloud_tpu dan jenis mesin Compute Engine untuk VM master. Anda dapat menggunakan metode ini untuk mengakses akselerator TPU v2 atau TPU v3. Akselerator TPU v3 tersedia dalam versi beta.

Mesin dasar berkemampuan TPU

Tetapkan tingkat skala ke BASIC_TPU untuk mendapatkan VM master dan VM TPU termasuk satu TPU dengan delapan core TPU v2, seperti yang Anda lakukan saat menjalankan contoh sebelumnya.

TPU worker dalam konfigurasi jenis mesin lama

Atau, Anda dapat menyiapkan konfigurasi mesin kustom jika memerlukan lebih banyak resource komputasi pada VM master:

• Tetapkan tingkat skala ke CUSTOM.
• Konfigurasi VM master untuk menggunakan jenis mesin lama yang sesuai dengan persyaratan tugas Anda.
• Tetapkan workerType ke cloud_tpu, untuk mendapatkan VM TPU yang mencakup satu Cloud TPU dengan delapan core TPU v2.
• Tetapkan workerCount ke 1.
• Jangan tentukan server parameter saat menggunakan Cloud TPU. Layanan menolak permintaan tugas jika parameterServerCount lebih besar dari nol.

Contoh berikut menunjukkan file config.yaml yang menggunakan jenis konfigurasi ini:

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

TPU worker dalam konfigurasi jenis mesin Compute Engine

Anda juga dapat menyiapkan konfigurasi mesin kustom dengan jenis mesin Compute Engine untuk VM master dan acceleratorConfig yang terpasang ke TPU VM.

Anda dapat menggunakan jenis konfigurasi ini untuk menyiapkan pekerja TPU dengan delapan core TPU v2 (serupa dengan konfigurasi tanpa acceleratorConfig) atau pekerja TPU dengan delapan core TPU v3 (beta). Baca selengkapnya tentang perbedaan antara akselerator TPU v2 dan TPU v3.

Penggunaan jenis mesin Compute Engine juga memberikan fleksibilitas yang lebih besar untuk mengonfigurasi VM master:

• Tetapkan tingkat skala ke CUSTOM.
• Konfigurasikan VM master untuk menggunakan jenis mesin Compute Engine yang sesuai dengan persyaratan tugas Anda.
• Tetapkan workerType ke cloud_tpu.
• Tambahkan workerConfig dengan kolom acceleratorConfig. Di dalam acceleratorConfig tersebut, tetapkan type ke TPU_V2 atau TPU_V3 dan count ke 8. Anda tidak boleh memasang core TPU dalam jumlah lain.
• Tetapkan workerCount ke 1.
• Jangan tentukan server parameter saat menggunakan Cloud TPU. Layanan menolak permintaan tugas jika parameterServerCount lebih besar dari nol.

Contoh berikut menunjukkan file config.yaml yang menggunakan jenis konfigurasi ini:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

Menggunakan Pod TPU

Pratinjau

Fitur ini tunduk pada "Persyaratan Penawaran Pra-GA" di bagian Persyaratan Layanan Umum dalam Persyaratan Khusus Layanan. Fitur pra-GA tersedia "sebagaimana adanya" dan mungkin memiliki dukungan terbatas. Untuk mengetahui informasi selengkapnya, lihat deskripsi tahap peluncuran.

Pod TPU adalah kumpulan perangkat TPU yang terhubung melalui antarmuka jaringan khusus berkecepatan tinggi. Pod TPU dapat memiliki hingga 2.048 core TPU, sehingga Anda dapat mendistribusikan beban pemrosesan di beberapa TPU.

Untuk menggunakan Pod TPU, Anda harus mengajukan permintaan penambahan kuota terlebih dahulu.

Contoh file config.yaml berikut menunjukkan cara menggunakan Pod TPU:

Pod TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

Pod TPU v3

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

Ada batasan jumlah core Pod yang dapat digunakan untuk setiap jenis TPU. Konfigurasi yang tersedia:

Jenis Pod TPU	Jumlah core Pod yang tersedia untuk digunakan
TPU_V2_POD	32, 128, 256, 512
TPU_V3_POD	32, 128, 256

Untuk mengetahui detail selengkapnya tentang cara memanfaatkan inti Pod TPU sepenuhnya, lihat dokumentasi Cloud TPU tentang Pod TPU.

Menggunakan container PyTorch yang telah dibangun sebelumnya pada TPU worker

Jika ingin melakukan pelatihan PyTorch dengan TPU, Anda harus menentukan kolom tpuTfVersion di trainingInput tugas pelatihan. Setel tpuTfVersion agar cocok dengan versi penampung PyTorch bawaan yang Anda gunakan untuk pelatihan.

AI Platform Training mendukung pelatihan dengan TPU untuk container PyTorch bawaan berikut:

URI image container	tpuTfVersion
gcr.io/cloud-ml-public/training/pytorch-xla.1-11	pytorch-1.11
gcr.io/cloud-ml-public/training/pytorch-xla.1-10	pytorch-1.10
gcr.io/cloud-ml-public/training/pytorch-xla.1-9	pytorch-1.9
gcr.io/cloud-ml-public/training/pytorch-xla.1-7	pytorch-1.7
gcr.io/cloud-ml-public/training/pytorch-xla.1-6	pytorch-1.6

Misalnya, untuk berlatih menggunakan container yang telah dibangun PyTorch 1.11, Anda dapat menggunakan file config.yaml berikut untuk mengonfigurasi pelatihan:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

Lihat juga bagian sebelumnya di halaman ini tentang Menggunakan TPU dalam kode PyTorch.

Menggunakan container kustom pada TPU worker

Beta

Fitur ini tunduk pada "Persyaratan Penawaran Pra-GA" di bagian Persyaratan Layanan Umum dalam Persyaratan Khusus Layanan. Fitur pra-GA tersedia "sebagaimana adanya" dan mungkin memiliki dukungan terbatas. Untuk mengetahui informasi selengkapnya, lihat deskripsi tahap peluncuran.

Jika ingin menjalankan container kustom di TPU worker, bukan menggunakan salah satu versi runtime Pelatihan AI Platform yang mendukung TPU, Anda harus menentukan kolom konfigurasi tambahan saat mengirimkan tugas pelatihan. Tetapkan tpuTfVersion ke versi runtime yang menyertakan versi TensorFlow yang digunakan penampung Anda. Anda harus menentukan versi runtime yang saat ini didukung untuk pelatihan dengan TPU.

Karena Anda mengonfigurasi tugas untuk menggunakan container kustom, AI Platform Training tidak menggunakan lingkungan versi runtime ini saat menjalankan tugas pelatihan Anda. Namun, Pelatihan AI Platform memerlukan kolom ini agar dapat menyiapkan pekerja TPU dengan benar untuk versi TensorFlow yang digunakan container kustom Anda.

Contoh berikut menunjukkan file config.yaml dengan konfigurasi TPU yang serupa dengan yang ada di bagian sebelumnya, kecuali dalam hal ini VM master dan pekerja TPU menjalankan container kustom yang berbeda:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V3
      count: 8

Jika menggunakan perintah gcloud beta ai-platform jobs submit training untuk mengirimkan tugas pelatihan, Anda dapat menentukan kolom API tpuTfVersion dengan flag --tpu-tf-version, bukan dalam file config.yaml.

Menggunakan TPUClusterResolver setelah TPU disediakan

Saat menggunakan container kustom, Anda harus menunggu hingga TPU disediakan sebelum dapat memanggil TPUClusterResolver untuk menggunakannya. Kode contoh berikut menunjukkan cara menangani logika TPUClusterResolver:

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

Pelajari lebih lanjut pelatihan terdistribusi dengan container kustom.

Langkah selanjutnya

• Pelajari lebih lanjut pelatihan model di Pelatihan AI Platform.
• Pelajari penyelarasan hyperparameter pada Pelatihan AI Platform, dengan memberi perhatian khusus pada detail untuk penyelarasan hyperparameter dengan Cloud TPU.
• Pelajari model referensi tambahan untuk Cloud TPU.
• Optimalkan model Anda untuk Cloud TPU dengan mengikuti praktik terbaik Cloud TPU.
• Lihat FAQ dan pemecahan masalah Cloud TPU untuk mendapatkan bantuan dalam mendiagnosis dan menyelesaikan masalah.