TensorFlow، PyTorch و JAX: انتخاب یک چارچوب یادگیری عمیق

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

• آیا باید از TensorFlow استفاده کنید؟
• آیا باید از PyTorch استفاده کنید؟
• آیا باید از JAX استفاده کنید؟
• TensorFlow در مقابل PyTorch در مقابل JAX

یادگیری عمیق هر روز زندگی ما را به روش های کوچک و بزرگ تغییر می دهد. خواه سیری یا الکسا باشد که دستورات صوتی ما را دنبال می کند، برنامه های ترجمه بلادرنگ روی تلفن های ما، یا فناوری بینایی رایانه ای که تراکتورهای هوشمند، ربات های انبار و ماشین های خودران را قادر می سازد، هر ماه به نظر می رسد پیشرفت های جدیدی را به ارمغان می آورد. و تقریباً همه این برنامه‌های یادگیری عمیق در یکی از سه چارچوب نوشته شده‌اند: TensorFlow، PyTorch و JAX.

از کدام یک از این چارچوب های یادگیری عمیق باید استفاده کنید؟ در این مقاله، نگاهی مقایسه ای در سطح بالا به TensorFlow، PyTorch و JAX خواهیم داشت. هدف ما این است که برخی از انواع برنامه‌هایی را که با نقاط قوت خود عمل می‌کنند، به شما ارائه دهیم، و همچنین عواملی مانند پشتیبانی جامعه و سهولت استفاده را در نظر بگیریم.

آیا باید از TensorFlow استفاده کنید؟

“هیچ کس به خاطر خرید IBM اخراج نشد” فریاد رایانش در دهه های ۱۹۷۰ و ۱۹۸۰ بود، و همین را می توان در مورد استفاده از TensorFlow در سال ۲۰۱۰ برای یادگیری عمیق. اما همانطور که همه ما می دانیم، IBM با ورود ما به دهه ۱۹۹۰ از بین رفت. آیا TensorFlow هنوز در این دهه جدید، هفت سال پس از انتشار اولیه آن در سال ۲۰۱۵، رقابتی است؟

خب، قطعا. اینطور نیست که TensorFlow در تمام این مدت ثابت مانده باشد. TensorFlow 1.x تماماً در مورد ساخت نمودارهای ایستا به روشی کاملاً غیر پایتون بود، اما با خط TensorFlow 2.x، همچنین می‌توانید مدل‌هایی را با استفاده از حالت “اشتیاقی” برای ارزیابی فوری عملیات بسازید و باعث می‌شود چیزها احساس بسیار بیشتری داشته باشند. مانند PyTorch. در سطح بالا، TensorFlow Keras را برای توسعه آسان‌تر به شما می‌دهد، و در سطح پایین، XLA (جبر خطی تسریع‌شده) کامپایلر را برای سرعت بهینه‌سازی می‌کند. XLA برای افزایش عملکرد در پردازنده‌های گرافیکی معجزه می‌کند، و این روش اصلی استفاده از قدرت TPUهای Google (واحد پردازش تنسور) است که عملکردی بی‌نظیر را برای مدل‌های آموزشی در مقیاس‌های عظیم ارائه می‌کند.

سپس همه چیزهایی وجود دارد که TensorFlow سال هاست به خوبی انجام داده است. آیا نیاز به ارائه مدل ها به شیوه ای کاملاً تعریف شده و قابل تکرار در یک پلت فرم بالغ دارید؟ سرویس TensorFlow برای شما آماده است. آیا نیاز دارید که استقرار مدل های خود را برای وب، یا برای محاسبات کم مصرف مانند گوشی های هوشمند، یا برای دستگاه های با محدودیت منابع مانند چیزهای IoT هدف گذاری کنید؟ TensorFlow.js و TensorFlow Lite هر دو در این مرحله بسیار بالغ هستند. و بدیهی است که با توجه به اینکه Google هنوز ۱۰۰٪ از استقرار تولید خود را با استفاده از TensorFlow اجرا می کند، می توانید مطمئن باشید که TensorFlow می تواند مقیاس شما را مدیریت کند.

اما… خوب، کمبود انرژی در اطراف پروژه وجود داشته است که این روزها نادیده گرفتن آن کمی سخت است. ارتقاء از TensorFlow 1.x به TensorFlow 2.x، در یک کلام، وحشیانه بود. برخی از شرکت‌ها تلاش‌های لازم برای به‌روزرسانی کد خود را برای کارکرد صحیح روی نسخه اصلی جدید بررسی کردند و در عوض تصمیم گرفتند کد خود را به PyTorch منتقل کنند. TensorFlow همچنین در جامعه تحقیقاتی که انعطاف پذیری PyTorch را چند سال پیش ترجیح می دادند، از دست داد و منجر به کاهش استفاده از TensorFlow در مقالات تحقیقاتی شد.

موضوع کراس نیز کمکی نکرده است. Keras دو سال پیش به بخشی از نسخه‌های TensorFlow تبدیل شد، اما اخیراً بار دیگر در یک کتابخانه جداگانه با برنامه انتشار خاص خود خارج شد. مطمئناً، جدا کردن Keras چیزی نیست که بر زندگی روزمره یک توسعه‌دهنده تأثیر بگذارد، اما چنین تغییر سطح بالایی در یک تجدید نظر جزئی در چارچوب، اعتماد به نفس ایجاد نمی‌کند.

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

آیا باید از PyTorch استفاده کنید؟

دیگر به پاشنه‌های تنسورفلو علاقه‌مند نیست، PyTorch یک نیروی اصلی در دنیای یادگیری عمیق امروزی است، شاید در درجه اول برای تحقیق، بلکه در برنامه های تولید بیشتر و بیشتر. و با تبدیل شدن حالت مشتاق به روش پیش‌فرض توسعه در TensorFlow و همچنین PyTorch، به نظر می‌رسد رویکرد Pythonic‌تر ارائه شده توسط تمایز خودکار PyTorch (autograd) در جنگ علیه نمودارهای استاتیک پیروز شده است.

برخلاف TensorFlow، PyTorch از زمان منسوخ شدن Variable API در نسخه ۰.۴. (قبلاً، Variable برای استفاده از autograd با تانسورها لازم بود؛ اکنون همه چیز یک تانسور است.) اما این بدان معنا نیست که چند قدم اشتباه در اینجا و آنجا وجود نداشته است. به عنوان مثال، اگر از PyTorch برای آموزش در چندین GPU استفاده می‌کنید، احتمالاً با تفاوت‌هایی بین DataParallel و DistributedDataParallel جدیدتر مواجه شده‌اید. تقریباً همیشه باید از DistributedDataParallel استفاده کنید، اما DataParallel در واقع منسوخ نشده است.

اگرچه PyTorch در پشتیبانی از XLA/TPU از TensorFlow و JAX عقب مانده است، اما از سال ۲۰۲۲ وضعیت بسیار بهبود یافته است. PyTorch اکنون از دسترسی به VM های TPU و همچنین سبک قدیمی تر پشتیبانی از TPU Node، همراه با فرمان آسان پشتیبانی می کند. استقرار خط برای اجرای کد شما بر روی CPU، GPU یا TPU بدون تغییر کد. و اگر نمی‌خواهید با برخی از کدهای دیگ بخاری که PyTorch اغلب شما را مجبور به نوشتن می‌کند مقابله کنید، می‌توانید به موارد اضافه شده در سطح بالاتر مانند PyTorch Lightning، که به شما امکان می دهد به جای بازنویسی حلقه های آموزشی، روی کار واقعی خود تمرکز کنید. از جنبه منفی، در حالی که کار بر روی PyTorch Mobile ادامه دارد، هنوز بسیار کمتر از TensorFlow Lite بالغ است.

از نظر تولید، PyTorch اکنون با پلتفرم‌های فریمورک آگنوستیک مانند Kubeflow ادغام شده است، در حالی که < یک پروژه href="https://pytorch.org/serve/" rel="nofollow">TorchServe می‌تواند جزئیات استقرار مانند مقیاس‌بندی، معیارها و استنتاج دسته‌ای را مدیریت کند و همه MLOps را در اختیار شما قرار دهد. /a> goodness در یک بسته کوچک که توسط خود توسعه دهندگان PyTorch نگهداری می شود. آیا PyTorch مقیاس می شود؟ متا سال‌هاست که PyTorch را در تولید اجرا می‌کند، بنابراین هر کسی که به شما بگوید PyTorch نمی‌تواند حجم کار را در مقیاس بزرگ انجام دهد، به شما دروغ می‌گوید. با این حال، موردی وجود دارد که ممکن است PyTorch به اندازه JAX برای دوره‌های آموزشی بسیار بسیار بزرگی که نیاز به بانک‌هایی از GPU یا TPU دارند، کاملاً دوستانه نباشد.

در نهایت، فیل در اتاق است. محبوبیت PyTorch در چند سال گذشته تقریباً به طور قطع با موفقیت کتابخانه Transformers Hugging Face مرتبط است. بله، Transformers اکنون از TensorFlow و JAX نیز پشتیبانی می کند، اما به عنوان یک پروژه PyTorch شروع شد و همچنان به چارچوب متصل است. با ظهور معماری ترانسفورماتور، انعطاف‌پذیری PyTorch برای تحقیق، و توانایی جذب بسیاری از مدل‌های جدید در عرض چند روز یا چند ساعت پس از انتشار از طریق مرکز مدل Hugging Face، به راحتی می‌توان فهمید که چرا PyTorch این روزها همه جا را فرا گرفته است. .

آیا باید از JAX استفاده کنید؟

اگر به TensorFlow علاقه ندارید، ممکن است Google چیز دیگری برای شما داشته باشد. به هر حال یه جورایی JAX یک چارچوب یادگیری عمیق است که توسط Google ساخته، نگهداری و استفاده می‌شود، اما به طور رسمی وجود ندارد. یک محصول گوگل با این حال، اگر به مقالات و انتشارات Google/DeepMind در یک سال گذشته نگاهی بیندازید، متوجه نمی شوید که بسیاری از تحقیقات گوگل به سمت JAX رفته است. بنابراین JAX یک محصول “رسمی” Google نیست، اما این چیزی است که محققان Google از آن برای عبور از مرزها استفاده می کنند.

JAX دقیقاً چیست؟ یک راه آسان برای فکر کردن در مورد JAX این است: یک نسخه تسریع‌شده با GPU/TPU از NumPy را تصور کنید که می‌تواند، با یک موج از یک عصا، به طور جادویی یک تابع پایتون را بردارید و تمام محاسبات مشتق شده در توابع گفته شده را انجام دهید. در نهایت، یک جزء JIT (Just-In-Time) دارد که کد شما را می گیرد و آن را برای کامپایلر XLA بهینه می کند و در نتیجه عملکرد قابل توجهی را نسبت به TensorFlow و PyTorch بهبود می بخشد. من شاهد افزایش سرعت اجرای برخی کدها به سادگی با پیاده‌سازی مجدد آن در JAX بدون انجام کار بهینه‌سازی واقعی، چهار یا پنج برابر شده‌ام.

با توجه به اینکه JAX در سطح NumPy کار می‌کند، کد JAX در سطح بسیار پایین‌تری نسبت به TensorFlow/Keras نوشته می‌شود، و بله، حتی PyTorch. خوشبختانه، یک اکوسیستم کوچک اما در حال رشد از پروژه‌های اطراف وجود دارد که بخش‌های اضافی را اضافه می‌کند. آیا کتابخانه های شبکه عصبی می خواهید؟ Flax از Google و هایکو از DeepMind (همچنین گوگل). Optax برای همه نیازهای بهینه ساز شما و PIX برای پردازش تصویر و موارد دیگر. هنگامی که با چیزی مانند Flax کار می کنید، ساخت شبکه های عصبی نسبتاً آسان می شود. فقط توجه داشته باشید که هنوز چند لبه ناهموار وجود دارد. به عنوان مثال، کهنه سربازان در مورد اینکه چگونه JAX اعداد تصادفی را متفاوت از بسیاری از چارچوب های دیگر مدیریت می کند، صحبت می کنند.

آیا باید همه چیز را به JAX تبدیل کنید و از آن پیشرفته‌تر استفاده کنید؟ خوب، شاید، اگر شما عمیقاً در تحقیقات مربوط به مدل‌های مقیاس بزرگ هستید که برای آموزش به منابع عظیمی نیاز دارند. پیشرفت‌هایی که JAX در زمینه‌هایی مانند آموزش قطعی و سایر موقعیت‌هایی که به هزاران غلاف TPU نیاز دارند، انجام می‌دهد، احتمالاً ارزش تغییر را دارند.

TensorFlow در مقابل PyTorch در مقابل JAX

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

با این حال، فکر نمی‌کنم این موضوع بحث‌انگیز باشد که بگوییم اگر در حوزه‌های متن و تصویر کار می‌کنید، و در حال انجام تحقیقات در مقیاس کوچک یا متوسط ​​با هدف به کارگیری این مدل‌ها در تولید هستید، پس PyTorch احتمالا بهترین گزینه شما در حال حاضر است. این روزها به نقطه شیرینی در آن فضا رسیده است.

با این حال، اگر نیاز دارید که هر ذره عملکرد را از دستگاه‌های کم‌محاسباتی کم کنید، آنگاه شما را به TensorFlow با بسته سنگی TensorFlow Lite هدایت می‌کنم. و در انتهای دیگر مقیاس، اگر روی مدل‌های آموزشی کار می‌کنید که در ده‌ها یا صدها میلیارد پارامتر یا بیشتر هستند، و عمدتاً آنها را برای اهداف تحقیقاتی آموزش می‌دهید، شاید وقت آن رسیده است که ارائه دهید. JAX یک چرخش.