Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

XLA: Mengoptimalkan Kompiler untuk Pembelajaran Mesin

XLA (Accelerated Linear Algebra) adalah kompiler khusus domain untuk aljabar linier yang dapat mempercepat model TensorFlow dengan kemungkinan tidak ada perubahan kode sumber.

Hasilnya adalah peningkatan dalam kecepatan dan penggunaan memori: sebagian besar benchmark internal berjalan ~ 1,15x lebih cepat setelah XLA diaktifkan. Dataset di bawah ini dievaluasi pada GPU NVidia V100 tunggal:

pengantar

Ketika program TensorFlow dijalankan, semua operasi dijalankan secara individual oleh pelaksana TensorFlow. Setiap operasi TensorFlow memiliki implementasi kernel GPU yang telah dikompilasi yang dikirim oleh pelaksana.

XLA menyediakan mode alternatif model yang sedang berjalan: ia mengkompilasi grafik TensorFlow ke dalam urutan kernel komputasi yang dihasilkan khusus untuk model yang diberikan. Karena kernel ini unik untuk model, mereka dapat mengeksploitasi informasi spesifik model untuk optimisasi. Sebagai contoh, mari kita lihat XLA optimasi yang dilakukan dalam konteks perhitungan TensorFlow sederhana:

 def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)
 

Berjalan tanpa XLA, grafik meluncurkan tiga kernel: satu untuk perkalian, satu untuk penambahan dan satu untuk reduksi. Namun, XLA dapat mengoptimalkan grafik sehingga menghitung hasilnya dalam satu peluncuran kernel. Ini dilakukan dengan "menggabungkan" penambahan, perkalian, dan pengurangan menjadi satu kernel GPU. Selain itu, operasi menyatu ini tidak menuliskan nilai-nilai perantara yang dihasilkan oleh y*z dan x+y*z ke memori; alih-alih "mengalirkan" hasil perhitungan menengah ini langsung ke penggunanya sambil tetap menyimpannya sepenuhnya dalam register GPU. Fusion adalah optimalisasi tunggal terpenting XLA. Bandwidth memori biasanya merupakan sumber daya yang paling langka pada akselerator perangkat keras, sehingga menghapus operasi memori adalah salah satu cara terbaik untuk meningkatkan kinerja.

Aktifkan XLA untuk model TensorFlow

Pengelompokan otomatis

Cara paling sederhana untuk mulai menggunakan XLA dalam model TensorFlow adalah mengaktifkan auto-clustering , yang secara otomatis menemukan cluster (subgraph yang terhubung) dalam grafik TensorFlow yang dapat dikompilasi dan dieksekusi menggunakan XLA. Pengelompokan otomatis pada GPU dapat diaktifkan dengan mengatur variabel lingkungan TF_XLA_FLAGS :

 $ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
 

Pengelompokan otomatis saat ini dioptimalkan untuk beban kerja GPU, tetapi juga dapat diaktifkan pada CPU dengan tambahan menggunakan flag --tf_xla_cpu_global_jit :

 $ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
 

Untuk contoh penggunaan rinci lihat colab tutorial auto-clustering .

Kompilasi eksplisit dengan fungsi tf.

Auto-clustering adalah alat yang hebat untuk membuat model lebih cepat tanpa ada perubahan pada kode, tetapi mungkin sulit untuk memahami perubahan apa yang telah dilakukan.

API kompilasi eksplisit menawarkan kontrol lebih halus untuk memilih fungsi mana yang harus dikompilasi. Misalnya, fungsi TensorFlow berikut yang melakukan pelatihan MNIST dikompilasi dengan XLA:

 @tf.function(experimental_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))
 

API experimental_compile memiliki semantik yang harus dikompilasi : seluruh fungsi dikompilasi dengan XLA, atau errors.InvalidArgumentError Pengecualian errors.InvalidArgumentError dilemparkan. XLA saat ini tidak dapat mengkompilasi fungsi di mana dimensi tidak dapat disimpulkan : yaitu, jika itu tidak mungkin untuk menyimpulkan dimensi dari semua tensor tanpa menjalankan seluruh perhitungan. Misalnya, fungsi berikut tidak akan dikompilasi:

 @tf.function
def not_compilable(x):
  return tf.unique(x)
 

Bentuk dapat bervariasi di seluruh proses:

 @tf.function(experimental_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
 

Lihat colab tutorial untuk contoh penggunaan yang lebih rinci.

Kompilasi AOT (Ahead-of-time) untuk CPU dengan tfcompile

Anda juga dapat menggunakan alat tfcompile mandiri, yang mengubah grafik TensorFlow menjadi kode yang dapat dieksekusi (hanya untuk CPU x86-64).

Periksa program yang dikompilasi

XLA menyediakan fasilitas introspeksi yang memungkinkan Anda memeriksa program yang dihasilkan. Untuk membuang program yang dihasilkan, gunakan variabel lingkungan XLA_FLAGS :

 $ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
 

Setelah dumping dilakukan, Anda dapat menemukan file-file berikut di /tmp/generated :

  • module_XXXX.*_optimizations.txt Program XLA yang dihasilkan, satu per setiap cluster yang dikompilasi. Melampirkannya saat mengirimkan laporan bug XLA sangat membantu!

  • module_XXXX.ir-*.ll File yang dihasilkan dalam representasi perantara LLVM , dengan NVPTX intrinsics.

  • module_XXXX.ptx File PTX yang dihasilkan.

Anda juga dapat membuang grafik yang memvisualisasikan penyertaan kluster XLA di dalam grafik TensorFlow dengan:

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
 

Laporan bug yang dapat direproduksi

Laporan bug jauh lebih mudah direproduksi jika mencakup dump untuk program XLA yang dihasilkan dan embedding auto-clustering yang digunakan. Untuk menghasilkan mereka untuk program TensorFlow yang berjalan dengan auto-clustering, luncurkan:

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"
 

Saat mengajukan bug, lampirkan konten /tmp/generated (dirujuk di atas).

Jika memungkinkan, cobalah untuk mengisolasi bug ke program XLA tunggal dengan menggunakan replay_computation dan menjalankannya secara iteratif pada program yang dihasilkan.

XLA Frontends

Terlepas dari TensorFlow, program XLA dapat dihasilkan oleh:

  • JAX : Transformasi komposable dari program Python + NumPy
  • Julia : Bahasa Julia untuk komputasi ilmiah
  • PyTorch : Kerangka kerja PyTorch

Bacaan lebih lanjut