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 tolok ukur internal bekerja ~ 1,15x lebih cepat setelah XLA diaktifkan. Dataset di bawah ini dievaluasi pada satu GPU NVidia V100:

pengantar

Saat program TensorFlow dijalankan, semua operasi dijalankan satu per satu oleh eksekutor TensorFlow. Setiap operasi TensorFlow memiliki implementasi kernel GPU yang telah dikompilasi yang dikirim oleh pelaksana.

XLA menyediakan mode alternatif untuk menjalankan model: XLA mengompilasi grafik TensorFlow menjadi urutan kernel komputasi yang dibuat secara khusus untuk model tertentu. Karena kernel ini unik untuk modelnya, kernel ini dapat memanfaatkan informasi khusus model untuk pengoptimalan. Misalnya, mari kita lihat pengoptimalan yang dilakukan XLA dalam konteks penghitungan 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 penjumlahan dan satu untuk pengurangan. Namun, XLA dapat mengoptimalkan grafik sehingga menghitung hasil dalam satu peluncuran kernel. Hal ini dilakukan dengan "menggabungkan" penambahan, perkalian dan reduksi menjadi satu kernel GPU. Selain itu, operasi fusi ini tidak menuliskan nilai antara yang dihasilkan oleh y*z dan x+y*z ke memori; alih-alih "mengalirkan" hasil dari penghitungan perantara ini langsung ke pengguna mereka sambil menyimpannya sepenuhnya dalam register GPU. Fusion adalah satu-satunya pengoptimalan terpenting XLA. Bandwidth memori biasanya merupakan sumber daya yang paling langka pada akselerator perangkat keras, jadi menghapus operasi memori adalah salah satu cara terbaik untuk meningkatkan kinerja.

Aktifkan XLA untuk model TensorFlow

Pengelompokan otomatis

Cara termudah untuk mulai menggunakan XLA dalam model TensorFlow adalah dengan mengaktifkan pengelompokan otomatis, yang secara otomatis menemukan cluster (subgraf yang terhubung) dalam grafik TensorFlow yang dapat dikompilasi dan dijalankan 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 menggunakan tanda tambahan --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 mendetail, lihat colab tutorial pengelompokan otomatis .

Kompilasi eksplisit dengan tf.fungsi

Pengelompokan otomatis adalah alat yang hebat untuk membuat model lebih cepat tanpa perubahan apa pun pada kode, tetapi mungkin sulit untuk memahami perubahan apa yang telah dilakukan.

API kompilasi eksplisit menawarkan kontrol yang lebih cermat untuk memilih fungsi mana yang harus dikompilasi. Misalnya, fungsi TensorFlow berikut yang menjalankan 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))

experimental_compile API memiliki semantik yang harus dikompilasi : apakah seluruh fungsi dikompilasi dengan XLA, atau terjadi errors.InvalidArgumentError Pengecualian errors.InvalidArgumentError dilempar. XLA saat ini tidak dapat mengompilasi fungsi yang dimensinya tidak dapat disimpulkan : yaitu, jika tidak memungkinkan untuk menyimpulkan dimensi semua tensor tanpa menjalankan seluruh komputasi. Misalnya, fungsi berikut tidak akan dikompilasi:

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

Bentuknya dapat bervariasi di semua jalur:

@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 detail.

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

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

Periksa program yang telah 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 berikut di /tmp/generated :

  • module_XXXX.*_optimizations.txt Menghasilkan program XLA , satu untuk setiap cluster yang dikompilasi. Melampirkan itu saat mengirimkan laporan bug XLA sangat membantu!

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

  • module_XXXX.ptx Membuat file PTX .

Anda juga dapat membuang grafik yang memvisualisasikan penyematan cluster 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 untuk direproduksi jika menyertakan dump untuk program XLA yang dihasilkan dan penyematan pengelompokan otomatis yang digunakan. Untuk membuatnya untuk program TensorFlow yang berjalan dengan pengelompokan otomatis, 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 melaporkan bug, lampirkan konten /tmp/generated (direferensikan di atas).

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

XLA Frontends

Selain TensorFlow, program XLA dapat dibuat oleh:

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

Bacaan lebih lanjut