‏حلقه‌ی آموزش PyTorch؛ جایی که ترتیب خط‌ها همه‌چیزه

این مقاله یه راهنمای دقیق برای نوشتن حلقه‌ی آموزش PyTorch ـه. حرف اصلی نویسنده اینه که ساختن این حلقه سخت نیست، ولی سر جای درست گذاشتن هر خط به‌طرز عجیبی شکننده‌ست: بعد از رفع خطاهای بدیهی، بقیه‌ی اشتباه‌ها رو به‌سختی می‌شه دید، چون هیچ‌کدوم اکسپشن پرت نمی‌کنن — فقط آموزش هم‌گرا نمی‌شه، نتیجه غلط درمیاد یا مصرف حافظه می‌ترکه. (آموزش توزیع‌شده و چند-GPU از دامنه‌ی این مقاله بیرونه.)

اسکلت کامل حلقه این شکلیه:

for epoch in range(100):
    model.train()
    for X_batch, y_batch in loader:
        optimiser.zero_grad()
        logits = model(X_batch)
        loss   = criterion(logits, y_batch)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimiser.step()
    scheduler.step()

مهم‌ترین بخش مقاله یه جدول «کجا ترتیب واقعاً مهمه» ـست. optimiser.zero_grad() اگه بعد از loss.backward() بیاد، گرادیان چند بچ روی هم انباشته می‌شه و آپدیت با مجموعشون انجام می‌شه، نه با بچ فعلی. clip_grad_norm_() اگه قبل از backward() صدا زده بشه، .grad هنوز خالیه و کلاً بی‌اثره؛ اگه بعد از optimiser.step() بیاد هم گرادیانی رو کلیپ می‌کنه که از قبل اعمال شده — باز هم بی‌اثر.

بقیه‌ی تله‌ها هم به همون اندازه بی‌صدان: scheduler.step() داخل حلقه‌ی بچ یعنی نرخ یادگیری به‌جای یک بار در هر epoch، به تعداد len(loader) بار افت می‌کنه. یادت رفت بعد از model.eval() دوباره model.train() بزنی؟ Dropout خاموش و BatchNorm فریز می‌مونه و مدل بی‌سروصدا تو حالت eval آموزش می‌بینه. torch.no_grad() رو موقع اعتبارسنجی جا انداختی؟ گراف autograd روی هر بچ ساخته می‌شه و حافظه تا OOM بالا می‌ره. حتی لاگ کردن خود loss به‌جای loss.item() گراف محاسباتی رو تو حافظه پین نگه می‌داره.

یه تله‌ی ظریف‌تر هم هست: model.to(device) بعد از ساختن optimizer. اگه هم‌زمان تبدیل نوع داده هم انجام بشه (مثل model.half())، متد to() آبجکت‌های nn.Parameter جدید می‌سازه و optimizer هنوز به پارامترهای دورریخته‌شده‌ی قبلی ارجاع می‌ده — یعنی آپدیت‌ها به تنسورهای اشتباهی اعمال می‌شن.

بخش دیتا هم چند نکته‌ی کاربردی داره. num_workers پروسه‌های جدا می‌سازه که بچ بعدی رو موازی با محاسبه‌ی GPU آماده می‌کنن؛ صفر یعنی پروسه‌ی اصلی همه‌ی لود رو انجام می‌ده و معمولاً GPU رو گشنه نگه می‌داره (دو تا چهار worker عملیه). pin_memory=True بچ رو تو حافظه‌ی pinned می‌ذاره تا موتور DMA کارت گرافیک مستقیم منتقلش کنه — فقط با num_workers > 0 و مقصد CUDA معنا داره. persistent_workers=True هم workerها رو بین epochها زنده نگه می‌داره تا هزینه‌ی fork دوباره نده.

drop_last=True هم بچ ناقص آخر رو دور می‌ندازه، چون آمار BatchNorm روی دو-سه نمونه نویزیه. درباره‌ی اندازه‌ی بچ، بچ کوچیک‌تر گرادیان نویزی‌تر می‌ده که خودش نوعی regularization ضمنیه، و بچ بزرگ‌تر حافظه‌ی بیشتری می‌خواد ولی موازی‌سازی بهتری می‌ده. توان‌های ۲ با تایل‌های tensor core (معمولاً ۱۶×۱۶ یا ۸×۱۶) هم‌تراز می‌شن، پس مضرب ۸ یا ۱۶ گرفتنِ ابعاد ایده‌ی خوبیه. آخرین توصیه هم تکرارپذیریه: seedها رو قبل از ساختن مدل و loader ست کنین.

نکات کلیدی:

• zero_grad() قبل از backward()؛ وگرنه گرادیان بچ‌ها روی هم انباشته می‌شه
• clip_grad_norm_() فقط بین backward() و step() معنا داره؛ جای دیگه بی‌اثره
• scheduler.step() بیرون از حلقه‌ی بچ، وگرنه نرخ یادگیری چند ده برابر سریع‌تر افت می‌کنه
• اعتبارسنجی رو داخل torch.no_grad() بذارین و بعدش حتماً model.train() رو برگردونین
• pin_memory فقط با num_workers > 0 و مقصد CUDA فایده داره؛ اندازه‌ی بچ رو مضرب ۸ یا ۱۶ بگیرین