Menggunakan TPU untuk melatih model

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

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

Menyiapkan lingkungan Google Cloud

Konfigurasikan lingkungan Google Cloud Anda dengan mempelajari bagian penyiapan di panduan memulai.

Memberi otorisasi Cloud TPU untuk mengakses project Anda

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

1. Dapatkan nama akun layanan Cloud TPU Anda 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 dalam project Anda, dengan peran Cloud ML Service Agent. Selesaikan langkah-langkah berikut di konsol Google Cloud atau menggunakan perintah gcloud:

Konsol

1. Login ke konsol Google Cloud dan pilih project tempat Anda menggunakan TPU.
2. Pilih IAM & Admin > IAM.
3. Klik tombol Tambahkan untuk menambahkan anggota ke project.
4. Masukkan akun layanan TPU di kotak teks Anggota.
5. Klik menu drop-down Roles.
6. Aktifkan peran Cloud ML Service Agent (Service Agents > Cloud ML Service Agent).

gcloud

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

   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 cara memberikan peran ke akun layanan, lihat dokumentasi IAM.

Contoh: Melatih contoh model MNIST

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

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

1. Download kode untuk model referensi TensorFlow dan buka direktori dengan 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. Tindakan ini memastikan bahwa perintah gcloud ai-platform jobs submit training menyertakan semua subpaket 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 yang dipaketkan ke bucket ini, dan AI Platform Training menyimpan output pelatihan di bucket.

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

Informasi selengkapnya tentang cara melatih model di Cloud TPU

Bagian ini memberi tahu Anda 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 region tempat TPU tersedia. Region berikut saat ini menyediakan akses ke TPU:

• us-central1
• europe-west4

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

Pembuatan versi TensorFlow dan AI Platform Training

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 selengkapnya tentang versi runtime Pelatihan AI Platform dan versi TensorFlow yang sesuai.

Kebijakan pembuatan versi sama dengan 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 gRPC TPU

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 melakukan satu perubahan penting saat menggunakan TPUClusterResolver untuk kode yang berjalan di AI Platform Training: Jangan berikan argumen apa pun saat Anda membuat instance TPUClusterResolver. Jika argumen kata kunci tpu, zone, dan project semuanya ditetapkan ke nilai default None, AI Platform Training akan otomatis memberikan resolver cluster dengan detail koneksi yang diperlukan melalui variabel lingkungan.

Contoh TensorFlow 2 berikut menunjukkan cara menginisialisasi resolver cluster dan strategi distribusi untuk pelatihan di 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 mesin, gunakan TPUStrategy API TensorFlow 2. Panduan TensorFlow untuk menggunakan TPU menunjukkan cara melakukannya.

Untuk melatih 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 bawaan, 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 melatih menggunakan TPU di AI Platform Training, Anda menggunakan satu perangkat XLA, bukan beberapa perangkat XLA.

Lihat juga bagian berikut di halaman ini tentang 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 mendorong server TensorFlow yang berjalan di pekerja TPU.

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

• 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 yang mendukung 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.

Pekerja TPU dalam konfigurasi jenis mesin lama

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

• Tetapkan tingkat skala ke CUSTOM.
• Konfigurasikan VM master untuk menggunakan jenis mesin lama yang sesuai dengan persyaratan tugas Anda.
• Tetapkan workerType ke cloud_tpu, untuk mendapatkan VM TPU yang menyertakan satu Cloud TPU dengan delapan core TPU v2.
• Tetapkan workerCount ke 1.
• Jangan menentukan 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

Pekerja TPU dalam konfigurasi jenis mesin Compute Engine

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

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

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

• 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 jumlah core TPU lainnya.
• Tetapkan workerCount ke 1.
• Jangan menentukan 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 oleh 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 core Pod TPU sepenuhnya, lihat dokumentasi Cloud TPU tentang Pod TPU.

Menggunakan container PyTorch bawaan di pekerja TPU

Jika ingin melakukan pelatihan PyTorch dengan TPU, Anda harus menentukan kolom tpuTfVersion di trainingInput tugas pelatihan. Tetapkan tpuTfVersion agar cocok dengan versi container 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 bawaan 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 penampung kustom di pekerja TPU

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 penampung kustom di pekerja TPU, bukan menggunakan salah satu versi runtime Pelatihan Platform AI 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 penampung kustom, AI Platform Training tidak menggunakan lingkungan versi runtime ini saat menjalankan tugas pelatihan Anda. Namun, AI Platform Training memerlukan kolom ini agar dapat menyiapkan pekerja TPU dengan benar untuk versi TensorFlow yang digunakan oleh penampung kustom Anda.

Contoh berikut menunjukkan file config.yaml dengan konfigurasi TPU yang mirip dengan yang ada di bagian sebelumnya, kecuali dalam hal ini VM master dan pekerja TPU masing-masing menjalankan penampung 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 penampung kustom, Anda harus menunggu 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 penampung kustom.

Langkah selanjutnya

• Pelajari lebih lanjut cara melatih model di AI Platform Training.
• Pelajari penyesuaian hyperparameter di Pelatihan AI Platform, dengan perhatian khusus pada detail untuk penyesuaian 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 bantuan mendiagnosis dan memecahkan masalah.