آشنایی با Mixed Precision و مزایای آن

دقت کنید که این الگوریتم فقط روی معماری‌های جدید Nvidia کار می‌کند. اگر مثل من بیشتر اوقات از Google Colab استفاده می‌کنید تنها وقتی می‌توانید از Mixed Precision استفاده کنید که T4 را به عنوان GPU در اختیار داشته باشید. ولی نگران نباشید، GPU های RTX هم از این قابلیت پشتیبانی می‌کنند. برای این که مطمئن شوید سخت‌افزار شما از این قابلیت پشتیبانی می‌کند باید بررسی کنید که Tensor Core دارد یا نه؛ چون Mixed Precision روی آن اجرا می‌شود. در صورتی که GPU شما Mixed Precision را پشتیبانی نکند، به دلیل رفت‌وآمدهایی که بین fp32 و fp16 انجام می‌دهد حتی باعث کاهش سرعت نیز می‌شود.

الگوریتم Mixed Precision چطور کار می‌کند؟

سوال بهتر این است که چه وقتی باید از fp32 استفاده کرد و چه وقتی از fp16؟ قبل از ابداع این روش هم تلاش‌های زیادی شده بود که شبکه‌های عصبی را فقط روی fp16 آموزش دهند؛ ولی مشکل این بود که به دلیل دقت کم‌تر شبکه‌هایی که به این روش آموزش می‌دیدند کسی از آن‌ها استفاده نمی‌کرد.

برای درک بهتر روش آموزش Mixed Precision نمودار فوق را درنظر بگیرید. در تصویر بالا - سمت چپ یک مدل معمولی را مشاهده می‌کنید که برای اجرای تمام فرآیندهایش از fp32 استفاده می‌کند. حال به تصویر دوم دقت کنید. توجه کنید که در این حالت ورودی هنوز به شکل fp32 باقی مانده است ولی فرایند forward روی fp16 انجام شده است که در واقع گام اول جهت افزایش سرعت است. همین‌طور مشاهده می‌کنید که برای محاسبه مقدار Loss آخرین خروجی شبکه عصبی تبدیل به fp32 شده است. علت این کار این است که مقدار Loss باید با بیشترین دقت ممکن محاسبه شود. یکی از دلایلی که مدل‌هایی که فقط از fp16 استفاده می‌کنند معمولا به دقت خوبی نمی‌رسند همین است که مقدار Loss را با تخمین بالایی محاسبه می‌کنند؛ اما در این روش چون مقدار Loss روی fp32 محاسبه می‌شود این مشکل پیش نخواهد آمد.

پس از محاسبه مقدار Loss مجددا آن را به fp16 تبدیل کرده و سپس فرآیند Backward انجام می‌شود که این کار نیز باعث افزایش سرعت می‌شود. در تصویر پایین - سمت چپ مشخص است که وزن‌ها را ابتدا با fp16 ‌ذخیره کرده؛ ولی بعد از محاسبه گرادیان آن‌ها را جهت به‌روزرسانی به fp32 تبدیل می‌کند. این کار به همان دلیلی برای Loss گفته شد انجام‌ می‌شود؛ در واقع گرادیان‌ها معمولا خودشان بسیار کوچک هستند و وقتی در fp16 اعمال شوند تقریبا باعث هیچ به‌روزرسانی‌ای روی وزن‌ها نمی‌شوند و شبکه آموزش داده نمی‌شود.

ممکن است کمی عجیب باشد؛ چون قبلا دیدیم که در فرآیند محاسبه یا Forward وزن‌ها به صورت fp16 بودند ولی برای ذخیره‌سازی و اعمال گرادیان‌ها از آن‌ها در حالت fp32 استفاده می‌شود. در حالت کلی فقط یک وزن وجود دارد که به آن Master Weights گفته می‌شود. این وزن‌ها همواره به صورت fp32 هستند و فقط زمانی که قرار است با آن‌ها محاسبات انجام دهیم (فرآیند Forward) به fp16 تبدیل می‌شوند.

چرا به Gradient Scaling نیاز داریم؟

اگر در حالت Forward برای یک لایه خاص ورودی float16 داشته باشید، Backward آن لایه نیز گرادیان‌ها را در float16 ایجاد می‌کند. مقادیر گرادیان‌ها به اندازه‌ای کوچک هستند که ممکن است در float16 قابل نمایش نباشند. این مقادیر به صفر میل می‌کنند (underflow) بنابراین به‌روزرسانی برای پارامترهای مربوطه از بین می‌رود و وزن‌ها هیچ تغییری نمی‌کنند. برای جلوگیری از نابود شدن این گردایان‌ها، Loss شبکه را در یک عدد بزرگ ضرب می‌کنیم و گرادیان‌ها را با استفاده از این Scaled Loss محاسبه می‌کنیم. پس از محاسبه گرادیان‌ها آن‌ها را به همان ضریبی که در Loss ضرب کرده بودیم تقسیم می‌کنیم. این کار باعث می‌شود تاثیر عملیاتی که روی Loss انجام دادیم از بین برود و روی آموزش مدل تاثیری نداشته باشد اما مشکل از بین رفتن گردایان‌ها به خاطر کوچک بودن مقادیرشان حل شده است. این فرآیند در سمت راست پایین تصویر نمایش داده شده است.

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

loss = model(inputs) # step one
# We assume gradients are float32. We do not want to divide float16 gradients
grads = compute_gradient(loss*512, model.weights)
grads /= 512 # step two# then update the weights

انتخاب یک عدد مناسب برای Scaling ممکن است کمی سخت باشد. اگر این عدد بیش از حد کوچک باشد مشکل از بین رفتن گرادیان‌های کوچک حل نشده باقی می‌ماند. هم‌چنین اگر خیلی بزرگ باشد مشکل برعکس شده و بسیاری از گرادیان‌ها مقدار بی‌نهایت خواهند داشت. البته جای نگرانی نیست چون PyTorch و Tensorflow به صورت خودکار این عدد را بسته به اندازه گرادیان‌های شبکه انتخاب می‌کنند.

پیاده‌سازی

تمام کدها همراه با خروجی در این notebook نیز در دسترس است.

پس از آشنایی با نحوه کار Mixed Precision استفاده از این تکنیک را در هر دو فریمورک اصلی یادگیری عمیق خواهیم دید. در این دو مثال ما یک مدل را با استفاده از Mixed Precision روی دیتاست CIFAR10 آموزش می‌دهیم.

PyTorch

از نسخه 1.6 به بعد Mixed Precision به PyTorch اضافه شده و می‌توان به راحتی از آن استفاده کرد. ابتدا کتابخانه‌های معمول را وارد می‌کنیم.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torchvision.models import mobilenet_v2

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

ttransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0), (255))])
train_ds = datasets.CIFAR10('./', download=True, transform=transform)
train_dl = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True)

برای استفاده از MP نیاز به وارد کردن amp داریم که به صورت خودکار فرایند تبدیل وزن ها را انجام می دهد.

# amp : Automatic Mixed Precision
from torch.cuda.amp import GradScaler # for gradient and loss scale
from torch.cuda.amp import autocast # Casts operations in float16 & 32 automatically

مثل همیشه یک Loss ،Model و Optimizer تعریف می‌کنیم.

model = mobilenet_v2()
model.classifier = nn.Linear(1280, 10)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=0.005)    
loss_fn = nn.CrossEntropyLoss().to(device)    

حال فقط کافی‌ست فرایند آموزش را تحت Auto Cast Context انجام دهیم. در واقع این Auto Cast تمام تبدیل‌ها را به صورت خودکار برای ما انجام می‌دهد. اگر پارامتر fp16 را برابر False قرار دهیم هیچ فرقی با یک آموزش معمولی نخواهد کرد.

def train(fp16=True, device='cuda'):
    scaler = GradScaler(enabled=fp16)
    loss_avg = 0.0
    for i, (inputs, labels) in enumerate(train_dl): 
      optimizer.zero_grad()
      # Casts operations to mixed precision
      with autocast(enabled=fp16):
          outputs = model(inputs.to(device))
          loss = loss_fn(outputs, labels.to(device))
          loss_avg = (loss_avg * i + loss.item()) / (i+1)
      # Scales the loss, and calls backward()
      # to create scaled gradients
      scaler.scale(loss).backward()
      # Unscales gradients and calls
      # or skips optimizer.step()
      scaler.step(optimizer)
      scaler.update()  
      if i0==0:   
          print('[%d, M] loss: %.4f' %(i, len(train_dl), loss_avg))

نکته‌ی دیگر این است که برای Loss و Optimizer از شئ Scaler استفاده کنید تا فرآیند Scaling را به‌طور خودکار انجام دهد.

مقدار خروجی‌ها در صورتی که از MP استفاده کنیم را می‌توانید در زیر ببینید:

train(fp16=True)
# outputs
[0,  391]      loss: 1.2953
[100,  391] loss: 1.2431
[200,  391] loss: 1.2172
[300,  391] loss: 1.2056

هم‌چنین مقدار خروجی‌ها در حالتی که مثل قبل شبکه را آموزش دهیم هم به شکل زیر است:

train(fp16=False)
# outputs
[0,  391]      loss: 1.2830
[100,  391] loss: 1.2331
[200,  391] loss: 1.2164
[300,  391] loss: 1.2011

مشخص است که از نظر کیفیت تفاوتی بین دو مدل نیست. برای اطلاعات بیشتر می توانید PyTorch Doc on Mixed Precision را بررسی کنید.

Tensorflow 2.X

ابتدا کتابخانه‌های مورد نیاز را وارد می‌کنیم.

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, Activation

برای استفاده از MP در این فریمورک نیاز دارید که یک سیاست سراسری برای مقدار دهی به لایه‌های مدل ایجاد کنید. با این کار تمام لایه‌های شبکه‌ی شما به طور پیش‌فرض از MP استفاده خواهند کرد. با این روش شما حتی می‌توانید روی TPU هم از Mixed Precision استفاده کنید.

from tensorflow.keras import mixed_precision
# set global dtype for all keras.layers
mixed_precision.set_global_policy('mixed_float16') # default is float32, if you use TPUs change it to mixed_bfloat16

همان‌طور که مشاهده می‌کنید تمام محاسبات روی fp16 صورت می‌گیرد ولی وزن‌های شبکه همان‌طور که قبل‌تر گفته شد روی fp32 ذخیره می‌شوند.

print('Compute dtype: ', mixed_precision.global_policy().compute_dtype)
print('Variable dtype: ', mixed_precision.global_policy().variable_dtype)
# outputs
Compute dtype:  float16
Variable dtype:  float32

در روند آماده کردن دیتا تفاوتی ایجاد نمی‌شود. درواقع Keras به نوع ورودی شما اهمیتی نمی‌دهد و شما می‌توانید مثل قبل دیتاست خود را بارگذاری کنید.

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

هر یک از لایه‌های Keras از سیاست سراسری برای ایجاد وزن‌ها استفاده می‌کنند مگر این که به‌طور صریح حالت دیگری مشخص شود. شما می‌توانید این ویژگی را با override کردن dtype تغییر دهید. نکته‌ی مهم این است که باید آخرین لایه‌ی شبکه عصبی را بدون توجه به نوع آن روی fp32 قرار دهیم تا بتوانیم Loss را روی fp32 محاسبه کنیم.

model = tf.keras.Sequential()
model.add(MobileNetV2(include_top=False, input_shape=(32, 32, 3))) 
model.add(Dense(10)) # use global policy which is float16 
# If your model ends in softmax, make sure it is float32. And regardless of what your model ends in, make sure the output is float32.
model.add(Activation('softmax', dtype='float32'))

پس فراموش نکنید که آخرین لایه هرچه که باشد (Dense یا Softmax یا هر لایه دیگر) باید dtype آن را برابر float32 قرار دهید تا شبکه بتواند Loss را با بالاترین دقت محاسبه کند.
حالا تمام قسمت‌ها را به یک تابع تبدیل می‌کنیم که مشخص می‌کند در روند آموزش از MP استفاده می‌شود یا خیر. بخش‌های دیگر هیچ تغییری نخواهند کرد. تابع Fit تمام کارهای دیگر نظیر Scaling را به طور خودکار برای ما انجام می‌دهد.

def train(fp16=True, epochs=1):
    # set floating point
    if fp16:
      mixed_precision.set_global_policy('mixed_float16')
    else:
      mixed_precision.set_global_policy('float32')
    # create & compile model
    model = tf.keras.Sequential()
    model.add(MobileNetV2(include_top=False, input_shape=(32, 32, 3)))
    model.add(Dense(10))
    model.add(Activation('softmax', dtype='float32')) # last layer must be float32
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam')
    # training 
    model.fit(x_train, y_train, epochs=epochs, batch_size=64)

اجرا با فعال کردن MP:

train(fp16=True)
782/782 [==============================] - 16s 52ms/step - loss: 1.6211

اجرا در حالت معمولی:

train(fp16=False)
782/782 [==============================] - 34s 37ms/step - loss: 1.6675

با مقایسه اجراهای فوق می‌توانید تفاوت سرعت و حتی کمی بهتر بودن کیفیت مدل را کنید.

استفاده از Mixed Precision در Custom Training Loop

شما می‌توانید از Mixed Precision در حالت Custom Training Loop هم استفاده کنید و در مدل‌های جدیدی که خودتان با استفاده از Keras Sub-classing ایجاد می‌کنید از مزایای MP بهره‌مند شوید. برای این کار نیاز دارید فرایند Gradient Scaling را درون Training Loop خود پیاده‌سازی کنید. راحت‌ترین روش برای انجام این کار این است که از کلاس LossScaleOptimizer استفاده کنید. Optimizer ای را که قبلا استفاده می‌کردید به عنوان ورودی به این کلاس بدهید و از این به بعد به جای آن، از شئ‌ای که این کلاس ایجاد می‌کند به عنوان Optimizer استفاده کنید. این کلاس دو مرحله به Optimizer معمولی شما اضافه می‌کند؛ یکی برای Loss Scaling و دیگری برای Gradient Scaling. برای درک بهتر این مطلب به مثال زیر توجه کنید.

class Fp16Training(tf.keras.Model):
    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
            # scale loss with optimizer
            scaled_loss = optimizer.get_scaled_loss(loss)
        # used scaled loss for compute gradient
        scaled_gradients = tape.gradient(scaled_loss, self.trainable_variables)
        # unscaled gradients to default value for stable training
        grads = optimizer.get_unscaled_gradients(scaled_gradients)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        # as usual
        self.compiled_metrics.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

مشاهده می‌کنید که از تابع get_scaled_loss برای Scale کردن Loss و از تابع get_unscaled_gradients برای Scale کردن گرادیان‌ها استفاده شده است. حال تنها نکته باقی‌مانده این است که از کلاس LossScaleOptimizer استفاده کنیم تا Optimizer ما آن دو تابع را در اختیار داشته باشد.

model = tf.keras.Sequential()
model.add(MobileNetV2(include_top=False, input_shape=(32, 32, 3)))
model.add(Dense(10))
# last layer or outputs must be float32 if use from_logits=True set dtype in last Dense
model.add(Activation('softmax', dtype='float32'))
# use custom trainig loop
cuistom_model = Fp16Training(model.inputs, model.outputs)
optimizer = keras.optimizers.Adam()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)
# compile model
cuistom_model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer)
cuistom_model.fit(x_train, y_train, batch_size=32, epochs=1)

برای اطلاعات بیشتر می‌توانید TF Doc on Mixed Precision را بررسی کنید.

سرعت بیشتر

معماری جدیدی که شرکت NVIDIA توسعه داده است می‌تواند به سرعت عملیات ماتریسی را در fp16 انجام دهد. این عملیات با استفاده از فناوری Tensor Core انجام می‌شود که بیشترین سرعت را زمانی دارد که سایز ماتریس‌های شما ضریبی از 8 باشد. تفاوتی نمی‌کند که از چه معماری شبکه‌ای عصبی استفاده می‌کنید. این فناوری با ماتریس‌هایی که اندازه‌ای از ضریب 8 دارند خیلی سریع‌تر کار می‌کند. نمونه کد زیر نحوه‌ای است که جهت ایجاد مدل جدید پیشنهاد می شود.

batch_size = 8*4
layer_input = 8*20
layer_output = 8*40
channel_number = 8*64