این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

XLA: کامپایلر بهینه سازی برای یادگیری ماشین

XLA (Algebra Linear Accelerated Linear) یک کامپایلر مخصوص دامنه برای جبر خطی است که می تواند مدلهای TensorFlow را با تغییر در کد منبع به طور بالقوه تسریع کند.

نتایج حاصل از پیشرفت در سرعت و استفاده از حافظه است: اکثر معیارهای داخلی بعد از فعال شدن XLA 1.15 برابر سریعتر عمل می کنند. مجموعه داده های زیر در یک پردازنده گرافیکی NVidia V100 واحد ارزیابی می شود:

مقدمه

هنگامی که یک برنامه TensorFlow اجرا می شود ، تمام عملیات به صورت جداگانه توسط مجری TensorFlow انجام می شود. هر عملیات TensorFlow دارای یک پیاده سازی هسته هسته GPU از پیش جمع شده است که مجری به آن اعزام می شود.

XLA یک حالت جایگزین از مدلهای در حال اجرا را فراهم می کند: نمودار TensorFlow را به دنباله ای از هسته های محاسبه که بطور خاص برای مدل داده شده تولید می کند ، می کشد. از آنجا که این هسته ها منحصر به فرد برای مدل هستند ، می توانند از اطلاعات خاص مدل برای بهینه سازی استفاده کنند. به عنوان مثال ، بیایید به یک بهینه سازی XLA در زمینه یک محاسبه ساده TensorFlow بپردازیم:

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

بدون XLA اجرا شود ، نمودار سه هسته را راه اندازی می کند: یکی برای ضرب ، دیگری برای اضافه کردن و دیگری برای کاهش. با این حال ، XLA می تواند نمودار را بهینه سازد تا نتیجه را در یک پرتاب هسته محاسبه کند. این کار را با "تلفیق" اضافی ، ضرب و کاهش به یک هسته GPU انجام می دهد. علاوه بر این ، این عملیات ذوب شده مقادیر میانی تولید شده توسط y*z و x+y*z به حافظه نمی نویسد. درعوض ، نتایج این محاسبات میانی را مستقیماً در اختیار کاربران خود قرار می دهد ، در حالی که آنها را به طور کامل در ثبات های GPU نگه می دارد. فیوژن مهمترین بهینه سازی XLA است. پهنای باند حافظه به طور معمول کمترین منبع برای شتاب دهنده های سخت افزاری است ، بنابراین از بین بردن عملیات حافظه یکی از بهترین راه ها برای بهبود عملکرد است.

XLA را برای مدلهای TensorFlow فعال کنید

خوشه بندی خودکار

ساده ترین روش برای شروع استفاده از XLA در مدل های TensorFlow ، فعال سازی خودکار خوشه بندی است ، که بطور خودکار خوشه ها (زیرگراف های متصل) را در نمودار TensorFlow پیدا می کند که می توان با استفاده از XLA کامپایل و اجرا شد. با تنظیم متغیر محیط TF_XLA_FLAGS می توان خوشه بندی خودکار در GPU را فعال کرد:

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

خودکار خوشه در حال حاضر برای حجم کار پردازنده گرافیکی بهینه سازی شده، اما آن را نیز می توانید بر روی CPU فعال باشد توسط علاوه بر این با استفاده از پرچم --tf_xla_cpu_global_jit :

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

برای یک مثال در مورد استفاده دقیق ، به کادر آموزش خودکار خوشه بندی مراجعه کنید .

تدوین صریح با عملکرد tf.funks

خوشه بندی خودکار ابزاری عالی برای سریعتر ساختن مدل بدون هیچ گونه تغییر در کد است ، اما درک این موضوع که تغییراتی که انجام شده است ممکن است کار سختی باشد.

API کامپایل صریح یک انتخاب دقیق تر را برای انتخاب عملکردهایی که باید تهیه شود ارائه می دهد. به عنوان مثال ، عملکرد زیر TensorFlow که آموزش MNIST را انجام می دهد با 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 دارای معنایی کاملاً کامپایل شده است : یا کل عملکرد با XLA وارد شده است ، یا یک errors.InvalidArgumentError استثناء errors.InvalidArgumentError پرتاب می شود. XLA در حال حاضر نمی تواند توابع را کاملاً کاملاً کاملاً کاملاً کاملاً کاملاً واضح سازد : یعنی در صورت عدم امکان استنباط ابعاد همه تنش ها بدون اجرای کل محاسبه امکان پذیر نیست. به عنوان مثال ، عملکرد زیر کامپایل نمی شود:

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

اشکال می تواند در هر دور متفاوت باشد:

 @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]))
 

برای مثال بیشتر در مورد استفاده ، به آموزش colab مراجعه کنید.

AOT (پیش از زمان) تلفیقی برای CPU با tfcompile

همچنین می توانید از یک ابزار tfcompile مستقل استفاده کنید ، که نمودار TensorFlow را به کد اجرایی تبدیل می کند (فقط برای CPU x86-64).

برنامه های کامپایل شده را بازرسی کنید

XLA امکانات بینش بینی را فراهم می کند که به شما امکان می دهد برنامه های تولید شده را بازرسی کنید. برای رها کردن برنامه های تولید شده ، از متغیر محیط XLA_FLAGS :

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

پس از انجام دامپینگ ، می توانید پرونده های زیر را در /tmp/generated :

  • module_XXXX.*_optimizations.txt برنامه های XLA ایجاد شده ، یک در هر خوشه گردآوری شده. پیوستن به آن هنگام ارسال گزارش اشکال XLA بسیار مفید است!

  • module_XXXX.ir-*.ll پرونده هایی را در بازنویسی LLVM و با ذات NVPTX ایجاد می کنید .

  • module_XXXX.ptx پرونده های PTX ایجاد شده است.

همچنین می توانید نمودار تجسم تعبیه خوشه های XLA را در داخل نمودار TensorFlow با:

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

گزارش اشکال قابل تکرار

اگر شامل زباله هایی برای برنامه های تولید شده XLA و تعبیه خوشه بندی خودکار مورد استفاده باشد ، گزارش اشکال بسیار ساده تر است. برای تولید آنها برای برنامه TensorFlow که با خوشه بندی خودکار اجرا می شود ، راه اندازی کنید:

 $ 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"
 

هنگام وارد کردن اشکالات ، محتویات فهرست /tmp/generated (در بالا ذکر شده) را پیوست کنید.

در صورت امکان ، با استفاده از replay_computation و اجرای replay_computation آن بر روی برنامه های تولید شده ، یک اشکال را به یک برنامه XLA اختصاص دهید.

XLA Frontends

به غیر از TensorFlow ، برنامه های XLA توسط:

  • JAX : تحولات قابل توجه برنامه های Python + NumPy
  • جولیا : زبان جولیا برای محاسبات علمی
  • PyTorch : چارچوب PyTorch

خواندن بیشتر