خواندن ۱۳ دقیقه·۲ سال پیش

آشنایی با Tensorflow XLA

یکی از نقاط قوت کتابخانه‌ی Tensorflow پویایی (dynamic) آن است. این ویژگی به محققان اجازه می‌دهد تا بتوانند ایده‌های خود را راحت‌تر و سریع‌تر بررسی کنند، بدون این که خود را درگیر ابعاد تک تک تنسورهای موجود در کد کنند. اما این ویژگی مثبت، باعث کاهش کارایی (performance) محاسبات تنسورفلو می‌شود. خیلی از اعمالی (operation) که توسط تنسورفلو انجام می‌شوند، به ابعاد داده‌ی ورودی وابسته‌ هستند و مشخص نبودن ابعاد تا زمان اجرا، باعث می‌شود که نتوان بسیاری از بهینه‌سازی‌ها را انجام داد، مانند بهینه‌سازی‌های مربوط به استفاده مجدد از حافظه که بدون دانستن حجم داده‌ها امکان‌پذیر نیست. Google در سال ۲۰۱۷، XLA را با این هدف معرفی کرد که راهکاری برای افزایش کارایی این کتابخانه با حفظ نقاط قوت کلیدی آن باشد.

مفاهیم مقدماتی

قبل از این که وارد مبحث XLA شویم، لازم است مروری بر برخی مفاهیم پایه‌ای کامپیوترها داشته باشیم.

زبان ماشین، اسمبلی و زبان‌های برنامه‌نویسی سطح بالا

کامپیوترها قادر به اجرای برنامه‌هایی هستند که به زبان ماشین نوشته شده باشند. زبان ماشین یک زبان سطح پایین است که از کدهای باینری تشکیل شده است. تمامی کدهایی که به هر کدام از زبان‌های برنامه‌نویسی نوشته شده باشند، برای اجرا باید در نهایت به زبان ماشین تبدیل شوند.

زبان اسمبلی یک زبان برنامه‌نویسی سطح پایین است که ارتباط قوی‌ای با زبان ماشین دارد. برنامه‌های نوشته‌شده به این زبان به آسانی توسط اسمبلرها به زبان ماشین تبدیل می‌شوند. هرچند کدهای این زبان برای انسان‌ها قابل فهم است، اما کسی که بخواهد به این زبان برنامه‌نویسی کند، باید با مفاهیم سخت‌افزاری آشنایی داشته باشد.

از این رو برنامه‌نویس‌ها معمولا ترجیح می‌دهند با زبان‌های سطح بالاتر که درکشان برای انسان‌ها راحت‌تر است کار کنند. این برنامه‌ها بعدا به روش‌های مختلف به اسمبلی یا زبان ماشین تبدیل و اجرا می‌شوند.

کامپایلر و مفسر

کامپایلرها و مفسرها هر دو وظیفه دارند تا برنامه‌های نوشته‌شده به زبان سطح بالا را به زبان ماشین تبدیل کنند. کامپایلرها این کار را قبل از اجرای برنامه انجام می‌دهند. آن‌ها کدهای نوشته‌شده در زبان برنامه‌نویسی سطح بالا را بررسی کرده و در صورت عدم وجود خطا آن را به یک فایل اجرایی با زبان ماشین تبدیل کرده و ذخیره می‌کنند. اما مفسرها این تبدیل را هم‌زمان با اجرای کد انجام می‌دهند. از این رو، مفسرها تنها خطوطی از برنامه را به زبان ماشین تبدیل می‌کنند که نوبت به اجرای آن‌ها برسد و لزوما کل برنامه را بررسی نمی‌کنند.

با توجه به این که کامپایلرها قبل از اجرا کد زبان ماشین را تولید کرده‌اند، اجرای آن‌ها نسبت به مفسرها سریع‌تر است. هم‌چنین چون کامپایلرها کل برنامه را قبل از تولید کد ماشین می‌بینند، می‌توانند بهینه‌سازی‌های مختلفی بر روی آن انجام دهند. اما به همین دلیل، لازم است تا تمامی نوع داده‌ها از قبل مشخص باشد و از این رو زبان‌های کامپایلری نمی‌توانند پویایی نوع داده‌ها (Dynamic typing) داشته باشند. C و C++ نمونه‌هایی از زبان‌های کامپایلری و python و PHP نمونه‌هایی از زبان‌های مفسری هستند.

معرفی XLA

مترجم XLA یا Accelerated Linear Algebra کامپایلری‌ست که به صورت اختصاصی برای محاسبات جبر خطی طراحی شده است و به برنامه‌نویسان اجازه می‌دهد بدون ایجاد تغییر اساسی در کدهای خود، سرعت اجرا را بالا ببرند. XLA با هدف افزایش سرعت اجرا، بهبود استفاده از حافظه و استفاده‌ی راحت‌تر از مدل روی موبایل‌ها و دیگر دستگاه‌ها توسعه داده شده است. همان‌طور که در تعریف کامپایلرها گفتیم، آن‌ها کدهای سطح بالا را به کد ماشین ترجمه می‌کنند. این اتفاق در XLA هم می‌افتد. به بیان دیگر، هر چیزی که بتوان آن را به گراف XLA تبدیل کرد، قابل ترجمه به زبان ماشین است.

برای این کار، XLA در دو مرحله کدها را بهینه می‌کند که می‌توان آن‌ها را در نمودار زیر دید.

زبان ورودی این کامپایلر XLA HLO یا High Level Operations است که می‌توان آن را به عنوان نمایش میانی دستورات در نظر گرفت. در مرحله اول، بهینه‌سازی‌هایی صورت می‌گیرند که به ماشین هدف وابسته نیستند و در مرحله بعد، بهینه‌سازی‌های دیگری با در نظر گرفتن ماشین مقصد انجام شده و سپس کد ماشین مربوطه تولید می‌شود.

یکی از مهم‌ترین بهینه‌سازی‌هایی که توسط XLA انجام می‌شود، ترکیب (fusion) است. به عنوان مثال، فرض کنید می‌خواهیم حاصل عبارت alpha*input0 + input1 را محاسبه کنیم. در حالت عادی و بدون XLA، برای محاسبه این عبارت دو کرنل اجرا می‌شوند، یکی برای ضرب و دیگری برای جمع. برای اجرای هر کدام از این کرنل‌ها، عملیات خواندن مقادیر ورودی از حافظه و ذخیره‌ی خروجی در آن تکرار می‌شود. اما می‌توان این دو عملیات را با هم ترکیب کرد. در این صورت با کاهش تعداد دفعات خواندن متغیرهای ورودی از حافظه و ذخیره نتیجه در آن، سرعت اجرای عملیات بالاتر می‌رود. به بیان دیگر، با استفاده از ترکیب یا fusion به جای دو تابع زیر، تابع سوم اجرا می‌شود.

در تصویر زیر که از این ارائه برداشته شده است، کرنل‌های اجرا شده برای محاسبه‌ی قطعه کد سمت چپ با استفاده از Tensorboard Profiler دیده می‌شوند. همان‌طور که می‌بینید، بخش زیادی از عملگرها با یکدیگر ترکیب شده و به عنوان fusion در یک کرنل اجرا شده‌اند.

مشاهده ترکیب کرنل‌ها در Tensorboard (منبع)

معرفی TF2XLA

تا کنون فهمیدیم که اگر یک گراف XLA داشته باشیم، می‌توانیم آن را کامپایل کرده و به کد ماشین‌ تبدیل کنیم، پس برای استفاده از XLA باید بتوانیم گراف تنسورفلو را به گراف XLA تبدیل کنیم. برای این که بفهمیم این تبدیل چگونه انجام می‌شود، بهتر است اول روند انجام محاسبات در تنسرفلو بدون استفاده از XLA را بررسی کنیم.

همان‌‌طور که در شکل فوق می‌بینید، کد تنسورفلویی که به زبان‌های مختلف نوشته شده است، در نهایت تبدیل به یک گراف تنسورفلویی می‌شود. در این گراف گره‌ها عملگر و یال‌ها عملوندها هستند و جهت آن‌ها جریان محاسبات را نشان می‌دهد. پس از ساخته شدن این گراف، مجری (Executer) گراف را پیمایش کرده و به دنبال گرهی می‌گردد که قابل محاسبه باشد و ورودی‌هایش آماده باشند. سپس، با توجه به عملگر و اطلاعات ماشین، از بین کرنل‌های مختلفی که به زبان C++ پیاده‌سازی شده‌اند، سریع‌ترین را انتخاب کرده و محاسبات را انجام می‌دهد.

حال نمودار زیر را در نظر بگیرید که روند اجرا در زمان استفاده از XLA را نشان می‌دهد. همان‌طور که می‌بینید، مجری این حالت با حالت قبل تفاوتی ندارد، اما مجموعه‌ی کرنل‌ها متفاوت هستند. این کرنل‌ها حاصل عملیات را محاسبه نمی‌کنند، بلکه وظیفه دارند گراف تنسورفلو را به گراف XLA تبدیل کنند.

تکرار این روند در نهایت بخش‌هایی از گراف اصلی را گراف‌های XLA تبدیل می‌کند که می‌توان آن‌ها را کامپایل کرد تا بتوان محاسبات مربوط به آن بخش را با سرعت بیشتری انجام داد. به این زیرگراف‌هایی که قابلیت تبدیل شدن به گراف XLA دارند خوشه یا cluster گفته می‌شود.

بکارگیری XLA در تنسورفلو

اکنون که می‌دانیم کدهای تنسورفلو چگونه توسط XLA به کدهای قابل اجرا توسط ماشین تبدیل می‌شوند، نگاهی به روش‌های مختلف استفاده از این قابلیت در تنسورفلو می‌اندازیم.

ترجمه‌ی JIT یا just in time

در این روش تنسورفلو ابتدا بخش‌هایی از گراف را که قابلیت تبدیل شدن به XLA دارند پیدا می‌کند. توجه داشته باشید که همیشه نمی‌توان انتظار داشت که کل گراف قابلیت تبدیل شدن به گراف XLA را داشته باشد؛ چون بعضی از کرنل‌های تنسورفلو به دلیل پویایی زیاد، قابل کامپایل نیستند. به علاوه، گاهی ممکن است کامپایل کردن یک خوشه باعث به هم ریختن نظم گراف اصلی شود. به عنوان مثال، کامپایل زیرگراف در شکل زیر باعث می‌شود که پیدا کردن ترتیب مناسب برای محاسبه گره‌ها غیر ممکن شود.

برای همین ممکن است مانند شکل زیر، تنها بخش‌هایی از گراف قابلیت بازنویسی به XLA را داشته باشند.

در این شرایط، مجری یا executer بعد از پیدا کردن گره مناسب اجرا، بررسی می‌کند که آیا آن گره یک cluster است یا نه و در صورتی که با یک خوشه مواجه باشد، زیرگراف را گراف XLA تبدیل کرده و سپس‌ آن را به زبان ماشین ترجمه می‌کند. هم‌زمان با این روند caching هم اتفاق می‌افتد، به این معنی که اگر این زیرگراف کامپایل‌شده را جای دیگری از گراف اصلی داشته باشیم، از کدی که قبلا به زبان ماشین تبدیل شده برای محاسبه نتیجه استفاده می‌کند. با درنظر گرفتن کارهای تکراری زیادی که در روند یادگیری ماشین اتفاق می‌افتد، می‌توان انتظار داشت که این تغییرات سرعت را بهبود بدهند.

ترجمه‌ی AoT یا ahead of time

در روش jit بخشی از گراف در زمان اجرا کامپایل می‌شد اما هدف AoT ترجمه کل گراف به صورتی است که مستقیما بر روی ماشین مقصد قابل اجرا باشد. برای این کار از ابزاری به نام tfcompile استفاده می‌شود. این ابزار، گراف تنسورفلو به همراه فایل تنظیماتی که در آن ورودی (feeds) و خروجی‌های (fetches) مدل مشخص می‌شوند را گرفته و با کمک XLA مدل را برای انواع ماشین‌های مختلف ترجمه می‌کند. خروجی tfcompile علاوه بر کدهای زبان ماشین، یک فایل C++ header نیز هست که چگونگی استفاده از مدل را مشخص می‌کند.

در نمونه کد زیر، می‌توانید نحوه استفاده از کد تولید شده را در زبان C++ ببینید. در این‌جا کتابخانه test_matmul به روش AoT کامپایل شده است.

نمونه کد

در این بخش برای درک بهتر چگونگی استفاده از XLA چندین نمونه کد را بررسی می‌کنیم.

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

در این بخش به توضیح آموزش طبقه‌بندی CIFAR-10 با کمک خوشه‌بندی خودکار (auto-clustering) در XLA می‌پردازیم. این آموزش از سایت رسمی Tensoflow برداشته شده است و نوت‌بوک آن از طریق این لینک قابل دسترسی است.

import tensorflow as tf

# Check that GPU is available: cf. https://colab.research.google.com/notebooks/gpu.ipynb
assert(tf.test.gpu_device_name())

tf.keras.backend.clear_session()
tf.config.optimizer.set_jit(False) # Start with XLA disabled.

ابتدا تنسورفلو را ایمپورت کرده و پس از اطمینان از دسترسی به GPU،‌ قابلیت jit را برای آزمایش اولیه غیرفعال می‌شود.

def load_data():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
    x_train = x_train.astype('float32') / 256
    x_test = x_test.astype('float32') / 256
    # Convert class vectors to binary class matrices.
    y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
    y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)
    return ((x_train, y_train), (x_test, y_test))

(x_train, y_train), (x_test, y_test) = load_data()

تابع load_data جهت آماده‌سازی داده‌ها نوشته شده است. این تابع پس از بارگیری داده‌ها از دیتاست‌های تنسورفلو، تصاویر را نرمال کرده و برچسب‌ها را به فرمت one-hot در می‌آورد.

def generate_model():
    return tf.keras.models.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), padding='same', input_shape=x_train.shape[1:]),
        tf.keras.layers.Activation('relu'),
        tf.keras.layers.Conv2D(32, (3, 3)),
        tf.keras.layers.Activation('relu'),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Dropout(0.25),
        tf.keras.layers.Conv2D(64, (3, 3), padding='same'),
        tf.keras.layers.Activation('relu'),
        tf.keras.layers.Conv2D(64, (3, 3)),
        tf.keras.layers.Activation('relu'),
        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
        tf.keras.layers.Dropout(0.25),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(512),
        tf.keras.layers.Activation('relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(10),
        tf.keras.layers.Activation('softmax')
    ])

model = generate_model()

یک مدل Sequential ساده برای آموزش طراحی شده و شئ model ساخته می‌شود.

def compile_model(model):
    opt = tf.keras.optimizers.RMSprop(lr=0.0001, decay=1e-6)
    model.compile(loss='categorical_crossentropy',
        optimizer=opt,
        metrics=['accuracy'])
    return model

model = compile_model(model)

مدل با RMSprop optimizer کامپایل می‌شود.

def train_model(model, x_train, y_train, x_test, y_test, epochs=25):
    model.fit(x_train, y_train, batch_size=256, epochs=epochs, validation_data=(x_test, y_test), shuffle=True)

def warmup(model, x_train, y_train, x_test, y_test):
    # Warm up the JIT, we do not wish to measure the compilation time.
    initial_weights = model.get_weights()
    train_model(model, x_train, y_train, x_test, y_test, epochs=1)
    model.set_weights(initial_weights)

warmup(model, x_train, y_train, x_test, y_test)
%time train_model(model, x_train, y_train, x_test, y_test)
scores = model.evaluate(x_test, y_test, verbose=1)

تابع train_model با هدف آموزش مدل در ۲۵ epoch و تابع warmup برای آماده کردن مدل جهت محاسبه‌ی زمان نوشته شده‌اند. تابع warmup، وزن‌های اولیه مدل را نگه داشته و پس از آموزش مدل به اندازه یک epoch،‌ آن‌ها را مجدد بر روی مدل قرار می‌دهد. با این کار، وزن‌های مدل تغییری نمی‌کنند اما خوشه‌های پیداشده توسط JIT کامپایل می‌شوند. سپس هر دوی این توابع فراخوانی می‌شوند. این روند، در نهایت مدل را به دقت ۰.۶۳۸ بر روی داده‌های تست می‌رساند و اجرای تابع train_model، به اندازه‌ی ۵۷.۱ ثانیه سخت‌افزار را درگیر می‌کند.

tf.keras.backend.clear_session()
tf.config.optimizer.set_jit(True) # Enable XLA.
model = compile_model(generate_model())
(x_train, y_train), (x_test, y_test) = load_data()
warmup(model, x_train, y_train, x_test, y_test)
%time train_model(model, x_train, y_train, x_test, y_test)

حال، می‌خواهیم تاثیر استفاده از JIT را بررسی کنیم. برای این کار ابتدا session را پاک کرده و سپس XLA JIT را فعال می‌کنیم. با این کار، زیرگراف‌هایی که قابلیت کامپایل شدن دارند، در اولین epoch ترجمه می‌شوند. سپس مدل را مجدد تعریف کرده و توابع warmup و train_model را مجدد فراخوانی می‌کنیم. می‌بینیم که این بار زمان آموزش مدل به ۴۹.۹ ثانیه کاهش پیدا کرده که می‌توان گفت سرعت آموزش مدل ۱.۱۴ برابر بهتر شده است.

اعداد گزارش‌شده در بالا از روی آموزش‌های تنسورفلو گفته شده‌اند. در تکرار این آزمایش توسط تیم شناسا بر روی Tesla T4 GPU، زمان آموزش بدون JIT برابر با ۱:۲۶ بود که پس از فعال‌سازی JIT به ۵۳.۱ ثانیه کاهش یافت که برابر با ۱.۶۲ برابر شدن سرعت آموزش است.

فعال‌سازی JIT با استفاده از tf.function و jit_compile=True

در مثال فوق، قابلیت autoclustering را فعال کردیم تا تمامی خوشه‌های ممکن به صورت خودکار پیدا شوند. به جای این کار، می‌توانیم XLA را تنها برای بخش‌های خاصی از کد فعال کنیم. در این بخش این آموزش از سایت Tensorflow را بررسی می‌کنیم که برای فعال‌سازی JIT از tf.function استفاده کرده است.

import tensorflow as tf
tf.compat.v1.enable_eager_execution()
# Size of each input image, 28 x 28 pixels
IMAGE_SIZE = 28 * 28
# Number of distinct number labels, [0..9]
NUM_CLASSES = 10
# Number of examples in each training batch (step)
TRAIN_BATCH_SIZE = 100
# Number of training steps to run
TRAIN_STEPS = 1000

ابتدا تنسورفلو را ایمپورت کرده و برخی مقادیر لازم مقداردهی می‌شوند.

# Loads MNIST dataset.
train, test = tf.keras.datasets.mnist.load_data()
train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()

# Casting from raw data to the required datatypes.
def cast(images, labels):
    images = tf.cast(tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)
    labels = tf.cast(labels, tf.int64)
    return (images, labels)

سپس دیتاست MNIST از مجموعه داده‌های تنسورفلو بارگیری شده و تغییرات لازم مانند تغییر نوع داده‌ها بر روی آن اعمال می‌شود.

layer = tf.keras.layers.Dense(NUM_CLASSES)
optimizer = tf.keras.optimizers.Adam()

در این مساله برای آموزش مدل از یک تک‌لایه‌ی Dense به همراه بهینه‌ساز Adam استفاده می‌شود.

@tf.function(jit_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))

تابع train_mnist برای آموزش مدل تعریف شده است. این تابع پس از گرفتن تصاویر و برچسب‌های مربوطه، عملیات back propagation را با استفاده از tf.GradientTape شبیه‌سازی می‌کند. توجه کنید که کل این تابع تحت tf.function decorator اجرا می‌شود. با پاس دادن jit_compile=True به این decoreator، قابلیت JIT فعال می‌شود.

for images, labels in train_ds:
    if optimizer.iterations > TRAIN_STEPS:
        break
    train_mnist(images, labels)

images, labels = cast(test[0], test[1])
predicted_labels = layer(images)
correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print(&quotPrediction accuracy after training: %s&quot % accuracy)

در نهایت با اجرای این کد مدل به اندازه‌ی ۱۰۰۰ گام آموزش داده می‌شود که این کار مدل را به دقت ۰.۸۷۲۱ بر روی داده‌های تست می‌رساند.

print(train_mnist.experimental_get_compiler_ir(images, labels)(stage='hlo'))

با اجرای تکه کد فوق می‌توانید دستورات کامپایل‌شده توسط XLA را ببینید. پیشنهاد می‌شود خروجی‌های این دستور را به ازای stage های مختلف (hlo, hlo_serialized, optimized_hlo, optimized_hlo_serialized, optimized_hlo_dot) و بر روی سخت‌افزارهای مختلف مشاهده کنید.

امیدواریم از خواندن مطلب استفاده کرده باشید. برای مطالعه بیشتر، پیشنهاد می‌شود این لینک را جهت آشنایی با مشکلات شایع در کار با XLA و این لینک را جهت استفاده از AoT بخوانید. هم‌چنین می‌توانید بخش پایانی ویدیوی موجود در این لینک را جهت تحلیل بهتر بهبودهای حاصل از استفاده از XLA مشاهده کنید.

کار بر روی Tensorflow XLA همچنان ادامه دارد. توسعه‌دهندگان تیم Tensorflow بر این باورند که قابلیت‌های این کامپایلر بیشتر از این‌هاست و امیدوارند که بتوانند در بهبودهای آتی فرآیند زمان‌بر آموزش مدل‌های سنگین را کوتاه‌تر کنند.

نویسنده: بهار برادران افتخاری

منابع:

tensorflow xla تنسورفلو performance کامپایلر

مجموعه دانش‌بنیان شناسا

شاید از این پست‌ها خوشتان بیاید