رندرینگ عصبی در NVIDIA OptiX با استفاده از بردارهای مشارکتی

انتشار NVIDIA OptiX 9.0 ویژگی جدیدی به نام بردارهای مشارکتی (cooperative vectors) را معرفی می‌کند که گردش‌کارهای هوش مصنوعی را به عنوان بخشی از هسته‌های ردیابی پرتو (ray tracing kernels) فعال می‌کند. این ویژگی از هسته‌های تنسور NVIDIA RTX برای عملیات ماتریسی شتاب‌یافته سخت‌افزاری و محاسبات شبکه عصبی در طول سایه‌زنی استفاده می‌کند. این امر تکنیک‌های رندرینگ هوش مصنوعی مانند NVIDIA RTX Neural Shaders و NVIDIA RTX Neural Texture Compression (NTC) را فعال می‌کند و گامی رو به جلو در جهت مواد فوتورئال با کیفیت فیلم در رندرینگ بی‌درنگ است.

رابط‌های برنامه‌نویسی (API) بردار مشارکتی در OptiX، DirectX، NVAPI، Slang و Vulkan معرفی شده‌اند. این مقاله به بررسی مفاهیم پشت بردارهای مشارکتی می‌پردازد که در تمام رابط‌های برنامه‌نویسی کاربرد دارند و یک مثال را با استفاده از رابط برنامه‌نویسی OptiX بررسی می‌کند.

یک پرسپترون چندلایه (multilayer perceptron یا MLP) بلوک سازنده اصلی بسیاری از الگوریتم‌های شبکه عصبی است. تحقیقات نشان داده‌اند که MLPها قادر به بازتولید دقیق اثراتی هستند که روی آن‌ها آموزش داده شده‌اند. حتی زمانی که به اندازه کافی کوچک باشند که در زمان واقعی اجرا شوند، MLPها قادر به مدیریت اثرات جالب مانند سایه‌زنی مبتنی بر فیزیک هستند، گاهی اوقات سریع‌تر و با ردپای حافظه کوچک‌تر از شبکه‌های سایه‌زنی سنتی.

MLP معمولاً از یک بردار ورودی، چندین لایه کاملاً متصل و یک بردار خروجی تشکیل شده است. اندازه بردارهای لایه‌های مختلف نباید یکسان باشد.

هر لایه از ارزیابی (استنتاج) MLP دارای دو فاز است: یک ترکیب خطی وزن‌دار و بایاس‌دار از مقادیر لایه قبلی و یک تابع فعال‌سازی (activation function) غیرخطی اختیاری. ترکیب خطی وزن‌دار و بایاس‌دار به یک ضرب ماتریس-بردار (matrix-vector multiply) و به دنبال آن اضافه کردن یک بردار بایاس تقلیل می‌یابد، که به عنوان یک تبدیل آفین (affine transformation) نیز شناخته می‌شود.

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

در یک MLP کاملاً متصل، هر نورون تابعی از تمام نورون‌های لایه قبلی است، همانطور که در شکل 1 نشان داده شده است. محاسبات مفهومی کامل یک لایه واحد در شکل 2 نشان داده شده است.

یکی از اهداف بردارهای مشارکتی فعال کردن استفاده از هسته‌های تنسور NVIDIA برای تسریع عملیات ماتریسی است. به طور معمول، مدل برنامه‌نویسی CUDA SIMT به یک پیچش کامل از رشته‌های فعال برای انجام این کار نیاز دارد، اما مدل برنامه‌نویسی ردیابی پرتو با رشته‌ها به طور مستقل رفتار می‌کند و یک پیچش کامل را تضمین نمی‌کند. علاوه بر این، هسته‌های تنسور ضرب ماتریس-ماتریس را ارائه می‌دهند، اما هر رشته ردیابی پرتو فقط به ضرب بردار-ماتریس نیاز دارد، که از هسته‌های تنسور به اندازه کافی استفاده نمی‌کند.

همچنین، رابط برنامه‌نویسی CUDA نیاز به هدف قرار دادن نسخه‌های سخت‌افزاری خاص دارد و تضمین نمی‌شود که از معماری به معماری دیگر سازگار باشد. برای آشنایی با رویکرد چندنخی CUDA برای ضرب ماتریس، راهنمای کاربر پس‌زمینه ضرب ماتریس را بررسی کنید.

بردارهای مشارکتی این محدودیت‌ها را با ارائه یک رابط برنامه‌نویسی که:

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

بردارهای مشارکتی می‌توانند از پس داده‌ها و واگرایی اجرا در یک پیچش برآیند، البته با کاهش عملکرد. بهترین عملکرد زمانی به دست می‌آید که وزن‌های MLP در یک پیچش یکسان باشند و زمانی که یک مکمل کامل از رشته‌ها در یک پیچش داشته باشید. استفاده از تغییر ترتیب اجرای سایه‌زن (Shader Execution Reordering یا SER) می‌تواند به دستیابی به هر دوی این اهداف کمک کند.

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

outputMatrix = inputMatrix × weightsMatrix + biasMatrix

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

بردار مشارکتی یک نوع بردار مات است، اساساً یک کلاس آرایه، که می‌تواند طول دلخواه داشته باشد. OptiX یک پیاده‌سازی از یک بردار مشارکتی به نام OptixCoopVec ارائه می‌دهد. بردارهای مشارکتی در OptiX از مجموعه خاص و محدودی از عملیات پشتیبانی می‌کنند که برای کمک به تسریع ارزیابی MLPها و شبکه‌های عصبی کوچک در نظر گرفته شده‌اند.

در رابط برنامه‌نویسی بردار مشارکتی، ضرب بردار-ماتریس با بایاس با تابع optixCoopVecMatMul انجام می‌شود که قسمت آفین ارزیابی لایه را انجام می‌دهد. از آنجایی که استفاده از توابع فعال‌سازی مختلف در مراحل مختلف اغلب مطلوب است، فعال‌سازی به طور جداگانه پس از ضرب بردار-ماتریس اعمال می‌شود و می‌تواند از مجموعه توابع برداری ارائه شده توسط رابط برنامه‌نویسی بردار مشارکتی ساخته شود.

outputVector = inputVector × matrix + bias

بردارهای مشارکتی در OptiX در همه دستگاه‌های RTX و GPUهای کلاس سرور انتخاب شده پشتیبانی می‌شوند. می‌توانید پشتیبانی دستگاه را با استفاده از optixDeviceContextGetProperty با OPTIX_DEVICE_PROPERTY_COOP_VEC پرس و جو کنید. استفاده از بردارهای مشارکتی در دستگاه‌های پشتیبانی نشده، خطا ایجاد می‌کند، زیرا هیچ پشتیبانی بازگشتی در دسترس نیست.

این بخش به بررسی رابط برنامه‌نویسی بردارهای مشارکتی برای انجام استنتاج یا ارزیابی لایه MLP در یک سایه‌زن OptiX می‌پردازد. مثالی که بررسی خواهیم کرد از نمونه optixNeuralTexture در OptiX SDK اقتباس شده است. این نمونه از NTC SDK استفاده می‌کند، که آموزش (فشرده‌سازی) بافت‌ها، ذخیره وزن‌ها و بایاس‌ها در یک قالب فایل مشخص شده و نشان دادن استنتاج (رفع فشرده‌سازی) در موقعیت‌های مختلف با زبان‌های سایه‌زنی مختلف را انجام می‌دهد. می‌توانید از NTC SDK برای فشرده‌سازی بافت‌های خود استفاده کنید و سپس آن‌ها را با استفاده از optixNeuralTexture مشاهده کنید.

کدی که از بردارهای مشارکتی استفاده می‌کند، به شدت از قالب‌بندی C++ استفاده می‌کند. قالب‌ها به کامپایلر کمک می‌کنند تا کد با کارایی بالا تولید کند با ارائه اندازه‌های آرایه استاتیک و انواع داده استاتیک که در زمان کامپایل شناخته شده‌اند. انواع پرکاربرد را از قبل تعریف کنید تا خواندن کد آسان‌تر شود.

به این ترتیب، می‌توانید از نوع میانبر T_OUT، برای مثال، به جای نوع OptixCoopVec<> قالبی استفاده کنید. تابع evalMLP یک MLP کامل را برای یک پیکسل معین روی صفحه ارزیابی می‌کند. از نظر شبه‌کد، ورودی‌ها را برای MLP تنظیم می‌کند، سپس هر لایه از MLP را ارزیابی می‌کند و در نهایت خروجی آخرین لایه را برمی‌گرداند:

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

ارزیابی لایه واحد کمی بیشتر از یک لفاف (wrapper) در اطراف optixCoopVecMatMul است. پس از ضرب ماتریس، افست را به بردار بایاس افزایش دهید تا آن را برای لایه بعدی آماده کنید (توجه داشته باشید که افست بایاس در این تابع با ارجاع ارسال می‌شود). سپس تابع فعال‌سازی را روی لایه فراخوانی کنید.

چیزی که ممکن است در این مثال‌های کد متوجه شده باشید این است که ما در حال ارسال اشاره‌گر پایه وزن یکسان به چندین فراخوانی به evalLayer هستیم و همچنین از این اشاره‌گر پایه برای وزن‌ها و بایاس‌ها استفاده می‌کنیم. داده‌های صحیح را در هر مرحله با افزودن مقادیر افست ثابت به اشاره‌گر پایه پیدا کنید، در این مورد یا weightsOffsetInBytes یا biasOffsetInBytes.

دو دلیل وجود دارد که کد به این روش نوشته می‌شود. دلیل اول این است که هنگام خواندن فایل‌ها در قالب NTC، API یک بلوک حافظه را برمی‌گرداند که در آن ماتریس‌های وزن و بردارهای بایاس همه به هم فشرده شده‌اند و می‌توانید از محاسبات ساده برای تکرار داده‌ها برای هر لایه استفاده کنید. دلیل دوم این است که این ساختار با آنچه که توسط API optixCoopVecMatMul انتظار می‌رود مطابقت دارد.

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

شمارش (enumeration) OPTIX_COOP_VEC_MATRIX_LAYOUT به شرح زیر تعریف شده است:

پارامتر طرح‌بندی داده یک راهنما برای کامپایلر در مورد قالب ماتریس است و الزامات هم‌ترازی حافظه و محدودیت‌های اندازه خاصی را تعیین می‌کند. ماتریس همیشه N ردیف و K ستون دارد. اولین کلمه در نام طرح‌بندی محور اصلی را تعیین می‌کند و کلمه دوم محور فرعی را تعیین می‌کند. هر کلمه محور یا "ردیف" یا "ستون" خواهد بود که به ردیف‌ها و ستون‌های ماتریس اشاره دارد.

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

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

پارامتر طرح‌بندی داده برای بردارهای مشارکتی در OptiX تنها از دو enum استفاده می‌کند: OPTIX_COOP_VEC_MATRIX_LAYOUT_COL_ROW_MAJOR و OPTIX_COOP_VEC_MATRIX_LAYOUT_ROW_COL_MAJOR. اگر می‌خواهید از یک ساختار داده ردیف-اصلی استفاده کنید، باید به عنوان OPTIX_COOP_VEC_MATRIX_LAYOUT_ROW_COL_MAJOR ذخیره شود. داده‌ها باید به ترتیب ردیف-ستون در حافظه مرتب شوند تا کامپایلر ضرب را به درستی برداری کند.

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

نمونه optixNeuralTexture در OptiX SDK نشان می‌دهد که چگونه می‌توان استنتاج بافت‌های عصبی RTX را با استفاده از بردارهای مشارکتی در یک سایه‌زن OptiX انجام داد. ویژگی‌های کلیدی که نمونه نشان می‌دهد عبارتند از:

• استفاده از بردارهای مشارکتی در سایه‌زن‌های OptiX برای انجام عملیات ماتریسی
• خواندن وزن‌ها و بایاس‌ها از یک قالب فایل باینری که توسط NTC SDK ایجاد شده است
• استفاده از قالب‌های C++ برای آسان‌تر کردن تولید کد با کارایی بالا توسط کامپایلر
• استفاده از ترکیبی از انواع داده و طرح‌بندی داده
• انتخاب یک تابع فعال‌سازی متفاوت در هر لایه

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

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

استفاده از بردارهای مشارکتی در OptiX می‌تواند به شما کمک کند تا ارزیابی MLPها و شبکه‌های عصبی کوچک را تسریع کنید و می‌تواند برای ایجاد مواد فوتورئال با کیفیت فیلم در زمان واقعی استفاده شود. برای کسب اطلاعات بیشتر، OptiX SDK و سایر APIهای ذکر شده در بالای این مقاله را ببینید.

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