PyTorch - מה זה בעצם @torch.inference_mode()?

תוכן עניינים

PyTorch - מה זה בעצם @torch.inference_mode()?

כשאנחנו מריצים מודל של PyTorch בזמן Inference (כלומר שלב שבו המודל רק מנבא, ולא לומד) - יש הרבה מנגנונים מיותרים שעדיין פועלים “ברקע”, כי PyTorch בנוי כברירת מחדל למצב אימון.

כאן נכנסת לתמונה ההוראה:

@torch.inference_mode()

מה היא עושה בפועל?

ה-decorator הזה מכבה לגמרי את מנגנון ה-Autograd, כלומר:

  • לא נשמרים גרפים של חישוב.
  • לא נאספת היסטוריה לצורך backpropagation.
  • לא מתבצעות פעולות מעקב אחרי gradients.

בקיצור: PyTorch עובר למצב קריאה בלבד - “אל תעקוב, רק תריץ”.

למה זה חשוב?

ב-inference, אנחנו לא צריכים לחשב גרדיאנטים - אז אין סיבה לשלם את מחיר הזיכרון והביצועים עליהם.

היתרונות העיקריים:

  • חיסכון של עד 30-50% בזיכרון GPU.
  • Latency נמוך יותר - פחות bookkeeping פנימי.
  • פחות עומס על CPU (שבו מנוהל ה-Autograd).

שימוש טיפוסי

import torch

model.eval()  # מכבה דברים כמו dropout
with torch.inference_mode():
    output = model(inputs)

או אם זו פונקציה שתרצי לעטוף:

@torch.inference_mode()
def predict(model, x):
    return model(x)

מה ההבדל בין inference_mode ל-no_grad?

שאלה קלאסית:

torch.no_grad() גם מכבה את מעקב הגרדיאנטים, אבל עדיין משאיר חלק ממבני הנתונים של Autograd זמינים.

torch.inference_mode() הולך רחוק יותר - הוא מניח שהטנזורים לא ישתנו, ולכן מבטל גם חלק מה-metadata. התוצאה: עוד יותר מהיר, במיוחד כשעובדים ב-batch גדול או בריצות חוזרות.

טיפ קטן

אם המודל שלך רץ באופן תדיר ב-production (API, שרת, benchmark) - העטיפה הזו היא חובה. היא לא משנה את התוצאה, אבל כן משנה דרמטית את היעילות.

סיכום

מצבמעקב גרדיאנטיםשינוי בטנזוריםחסכון בזיכרוןמתאים ל-
רגילאימון
no_gradבינוניבדיקה / ניסוי
inference_modeגבוהProduction / Deployment

תגובות