یادگیری توزیع شده در تنسورفلو Distributed Learning with TensorFlow

سلام این اولین پست منه تو ویرگول و می‌خوام از یه مسئله باحال تو تنسورفلو حرف بزنم. به صورت کلی وقتی مدل‌هاتون رو تو تنسورفلو دولوپ می‌کنید، تنسورفلو میره سراغ GPUدر منابع کلاینت یا سرور شما. خب حالا اگه ما تو سیستمی که داریم، چندتا گرافیک داشته باشیم چی؟ اگه چندتا سیستم داشته باشیم که هر کدوم یه دونه گرافیک داشته باشن چی؟ اگه چندتا سیستم که هر کدوم چندتا گرافیک داشته باشن چی؟ چیکار کنیم که همه گرافیک‌ها با هم و در کنار هم برای آموزش با هم همکاری کنند؟ و کلی سوال دیگه.

از یک مسئله راحت شروع می‌کنیم، یعنی چند GPU در یک سیستم. خب بیاید بریم تو کار ببینیم چه طوری‌هاست:

اول میایم از لایبرری distribute تنسورفلو MirroredStrategy رو صدا می‌کنیم، این سلطان خودش تشخیص می‌ده چندتا GPU سیستم داره:

MyStrategy = tf.distribute.MirroredStrategy()

حالا اگه تو یه سیستم ما سه تا GPU داشته باشیم ولی بخوایم از دوتاش استفاده کنیم:

MyStrategy = tf.distribute.MirroredStrategy(devices=[&quot/gpu:0&quot, &quot/gpu:1&quot])

برای اینکه بخوایم تعداد GPUهایی که قراره تو این یادگیری مشارکت داشته باشند رو با کد زیر می‌تونیم بفهمیم:

MyStrategy .num_replicas_in_sync

حالا به بحث مهمی می‌رسیم که چطوری این اتفاق منجر میشه تا یادگیری به صورت توزیع شده انجام بشه؟ خب ما چیزی تحت عنوان Batch Size داریم که با این پارامتر تعیین می‌کنیم تعداد ورودی داده‌ها به شبکه عصبی در دسته‌های چندتایی باشه. حالا این Batch Sizeرو در یادگیری توزیع شده اینطوری بدست می‌آریم:

batch_size_per_replica = 64
batch_size = batch_size_per_replica * MyStrategy .num_replicas_in_sync

مثلا اگه اندازه دسته‌های با یک کارت گرافیک 64 تا باشه با 3 تا کارت گرافیک می‌شه 192تا. این یعنی سریع‌تر شده یادگیری اونم حدودا سه برابر:))))))

حالا اگه بخوایم مدل رو تعریف کنیم چه کاری باید بکنیم؟

with MyStrategy .scope():
#مدل رو اینجا تعریف می‌کنیم

حالا اگه چندتا سیستم داشته باشیم چی؟ اولین خط کد‌ در بالا رو اینطوری باز نویسی می‌کنیم(جای MirroredStrategy می‌ذاریم MultiWorkerMirroredStrategy):

MyStrategy = tf.distribute.MultiWorkerMirroredStrategy()

خب حالا از کجا قراره بفهمه سیستم‌های ما کجاست و کدوم هستن؟ تنسورفلو از مفهومی به نام کلاستر استفاده می‌کنه. یه پارامتری وجود داره به اسم TF_CONFIG که شاید یه روزی در موردش کامل توضیح دادم ولی الان در این حد می‌گم که میفهمونه به تنسورفلو که کدوم دیوایس تو کدوم کلاستر بیوفته و worker چی باشه براش(هر ورکر مجموعه‌ای از چند جاب و هر جاب شامل چند تسکه). معرفی کردنش اینطوریه:

import json
import os
import sys
import tensorflow as tf
import mnist_setup

# اول همه کارت گرافیک‌ها رو غیر فعال می‌کنیم
os.environ[&quotCUDA_VISIBLE_DEVICES&quot] = &quot-1&quot

# TF_CONFIG ریست
os.environ.pop('TF_CONFIG', None)

#TF_CONFIG تنظیم
tf_config = {
    'cluster': {
        'worker': [&quothost1:port&quot, &quothost2:port&quot, &quothost3:port&quot]
     },
     'task': {'type': 'worker', 'index': 0}
}
json.dumps(tf_config)

#میریم سراغ استراتژی و مدل مثلا دیتاست ما مینسته
per_worker_batch_size = 64
tf_config = json.loads(os.environ['TF_CONFIG'])
num_workers = len(tf_config['cluster']['worker'])
MYstrategy = tf.distribute.MultiWorkerMirroredStrategy()
global_batch_size = per_worker_batch_size * num_workers
multi_worker_dataset = mnist_setup.mnist_dataset(global_batch_size)

with MYstrategy .scope():
    # تعریف مدل

به TF_CONFIGنگاه کنید، اگر به آن پارامتر سرور اضافه کنیم معماری تغییر می‌کنه که به این معماری اصطلاحا PS-Worker architecture گفته می‌شه. تو این معماری تو یک کلاستر ما دو نوع سرور داریم: سرورهای پارامتر(parameter server) و ورکرها. سرورهای پارامتر پارامترهای مدل رو ذخیره می کنند و ورکرها گرادیان‌ها رو محاسبه می کنند. در هر تکرار، ورکرها پارامترهایی رو از سرورهای پارامتر دریافت می‌کنند و بعدش گرادیان‌های محاسبه‌شده رو به سرورهای پارامتر برمی‌گردونن. سرورهای پارامتر، گرادیان های محاسبه شده رو جمع می‌کنند، پارامترها رو آپدیت کرده و پارامترهای آپدیت شده رو برای ورکرها توزیع می‌کنند. شکل زیر رو ببینید:

اگه شما پارامتر سرور رو وارد کنید استراتژی شما بر اساس معماری توضیح داده میشه و باید به جای استراتژی MultiWorkerMirroredStrategy از استراتژی ParameterServerStrategy استفاده کنید. کانفیگش اینطوری داده می‌شه:

os.environ[&quotTF_CONFIG&quot] = json.dumps({
    &quotcluster&quot: {
        &quotworker&quot: [&quothost1:port&quot, &quothost2:port&quot, &quothost3:port&quot],
        &quotps&quot: [&quothost4:port&quot, &quothost5:port&quot]
    },
    &quottask&quot: {&quottype&quot: &quotworker&quot, &quotindex&quot: 1}
})

خب حالا اگه شما یک سرور دارید با چند GPU و نمی‌خواهید هر بار وزن‌های بدست اومده در هرکدوم به صورت جداگانه ذخیره نشوند و دوباره کاری رخ ندهد. اینجا میشه رفت سراغ CentralStorageStrategy. این استراتژی وزن‌های بدست اومده تو هر مرحله رو داخل CPU ذخیره می‌کنه و تمامی پردازش‌ها همچنان به صورت سینک(پایین توضیحش می‌دم) روی باقی GPU‌ها صورت می‌گیره. این استراتژی مانند استراتژی ParameterServerStrategy از مدل فیت کراس و لوپ شخصی سازی شده آموزش، ساپورت رسمی نمی‌کنند اما به صورت اکسپریمنتال ساپورت می‌کنند که برای فرخوانی اونها باید اینطوری کد بزنید:

MYstrategy = tf.distribute.experimental.ParameterServerStrategy ()
#اگه استراتژی شما پارامتر سرور بود

MYstrategy = tf.distribute.experimental.CentralStorageStrategy()
#اگه استراتژی شما سنترال استوریج بود

خب حالا سینک و غیرسینک بودن در استراتژی‌ها یعنی چی؟

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

سینک(Synchronous) به این معنیه که همه ورکرها هماهنگ هستن با همدیگه. هر ورکر پارامترها رو می‌خوانه، یه گرادیان رو محاسبه می‌کنه و منتظر می‌مونه تا سایر ورکرها کارشون رو تموم کنند. سپس الگوریتم یادگیری میانگین همه گرادیان‌هایی رو که محاسبه کردن ورکرها رو محاسبه می کنه و یک میانگین رو آپدیت می‌کنه.

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

گفتیم که تنسورفلو توزیع شده از یک خوشه شامل یک یا چند سرور پارامتر و ورکر تشکیل شده. از اونجایی که ورکرها در طول آموزش گرادیان‌ها رو محاسبه می کنن، معمولاً روی یک GPU قرار می‌گیرند. سرورهای پارامتر فقط باید گرادیان‌ها را جمع آوری کنند و آپدیت‌ها رو پخش کنن، بنابراین معمولاً روی CPU قرار می گیرند نه GPU. یکی از ورکرها، بهش میگن ورکر چیف، مدل آموزشی رو داره و به بقیه می‌ده، اینشیالایز می‌کنه مدل رو، تعداد مراحل آموزشی تکمیل شده رو شمارش می‌کنه، مانیتور می‌کنه سشن‌ها رو، لاگ ها رو برای TensorBoard ذخیره می‌کنه و چک پوینت مدل رو ذخیره و بازیابی می‌کنه. ورکر چیف همچنین خرابی‌ها را مدیریت می کنه. اگر خودش از کار بیوفته، آموزش باید از آخرین چک پوینت دوباره شروع بشه.

• یکی از معایب Distributed TensorFlow، بخشی از هسته TensorFlow، اینه که شما باید راه اندازی و توقف سرورها رو دستی مدیریت کنید. به طور کلی، این منجر به تعداد زیادی دستور سوئیچ در کد شما می‌شه تا مشخص بشه کدام بخش از کد باید در سرور فعلی اجرا شوند. ClusterSpec مجموعه ای از فرآیندهایی رو نشون می‌ده که در محاسبات تنسورفلو توزیع شده شرکت می کنند. هر tf.distribute.Server در یک کلاستر خاص ساخته می‌شه. بیاید یه ClusterSpec ست کنیم:

def main(_):
  ps_hosts = FLAGS.ps_hosts.split(&quot,&quot)
  worker_hosts = FLAGS.worker_hosts.split(&quot,&quot)

  # یه کلاستر می‌سازیم
  cluster = tf.train.ClusterSpec({&quotps&quot: ps_hosts, &quotworker&quot: worker_hosts})

  # استارت می‌کنیم سرور رو برای تسک‌‌ها و جاب‌ها
  server = tf.train.Server(cluster,
                           job_name=FLAGS.job_name,
                           task_index=FLAGS.task_index)

with tf.device(&quot/job:ps/task:0&quot):
     weights_1 = tf.Variable(...)
     biases_1 = tf.Variable(...)

with tf.device(&quot/job:ps/task:1&quot):
     weights_2 = tf.Variable(...)
     biases_2 = tf.Variable(...)

with tf.device(&quot/job:worker/task:7&quot):
     input, labels = ...
     layer_1 = tf.nn.relu(tf.matmul(input, weights_1) + biases_1)
     logits = tf.nn.relu(tf.matmul(layer_1, weights_2) + biases_2)
  # ...

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

حرف آخر

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

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