no_grad, inference_mode, 분리

'기울기를 추적하지 않는다'고 말하는 세 가지 방법

학습할 때는 대부분 autograd 기록이 필요하지만 추론, 평가, 전처리, 평가지표 계산, 시각화에서는 필요하지 않아. 그래프를 기록하면 활성화를 보관해야 해서 메모리를 쓰고 CPU 비용도 조금 들어. 이 기록을 건너뛰는 도구는 세 가지야:

torch.no_grad(): 문맥 관리자야. 블록 안에서 실행한 연산은 그래프를 만들지 않아. model.eval()로 평가할 때 흔히 함께 사용해.
torch.inference_mode(): no_grad보다 강한 방식이야. autograd와 뷰 추적을 모두 비활성화해서 조금 더 빠르며, 순수 추론에는 이쪽을 권장해.
tensor.detach(): 저장 공간은 공유하지만 autograd 그래프에서 분리된 새 텐서를 반환해. 코드 블록 전체가 아니라 손실 기록이나 강화학습의 목표 신경망처럼 개별 텐서만 분리할 때 써.

경계: `model.eval()`과 `no_grad()`

둘은 관련이 있지만 역할이 달라:

model.eval()은 모듈의 동작을 바꿔. 드롭아웃은 비활성화하고 BatchNorm은 배치 통계 대신 누적 통계를 사용하지만, autograd 자체를 끄지는 않아.
torch.no_grad()는 autograd를 꺼. 모듈의 동작은 바꾸지 않아.

그래서 평가할 때는 거의 항상 둘 다 필요해. 올바른 계층 동작에는 model.eval(), 속도와 메모리 절약에는 torch.no_grad() 또는 inference_mode()를 써.

Code

no_grad와 inference_mode: 문맥 관리자·python

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
x = torch.randn(4, 10)

# no_grad: context manager
with torch.no_grad():
    y = model(x)
    print(y.requires_grad)        # False — no graph built

# inference_mode: stronger and faster
with torch.inference_mode():
    y = model(x)
    print(y.requires_grad)        # False

# As a decorator
@torch.inference_mode()
def predict(model, x):
    return model(x)

개별 텐서를 그래프에서 분리하기·python

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2

# Logging the value — don't drag autograd along
loss_value = y.detach().item()    # plain Python float

# A target network in RL — gradients should NOT flow into it
target = (x ** 2).detach()         # treated as a constant from now on

# detach IS view-shaped — same storage, different graph status
print(y.detach().data_ptr() == y.data_ptr())  # True

완전한 평가 기본 패턴·python

import torch
import torch.nn as nn

def evaluate(model, val_loader, criterion, device):
    model.eval()                       # behavior switch (dropout off, etc.)
    total_loss = total_correct = total_n = 0

    with torch.inference_mode():       # autograd switch (off)
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            out = model(x)
            loss = criterion(out, y)

            total_loss += loss.item() * x.size(0)
            total_correct += (out.argmax(-1) == y).sum().item()
            total_n += x.size(0)

    return total_loss / total_n, total_correct / total_n

Exercise

작은 모델을 하나 골라 추론 100회의 실행 시간을 세 조건에서 재 봐. (a) 별도 autograd 문맥 없이 실행, (b) torch.no_grad 사용, (c) torch.inference_mode 사용. 작은 모델에서도 보통 (a), (b), (c) 순으로 실행 시간이 줄어들 거야. 측정값을 보관해 두면 평가 시간의 허용 범위를 잡을 때 유용해.

'기울기를 추적하지 않는다'고 말하는 세 가지 방법

경계: `model.eval()`과 `no_grad()`

Code

External links

Exercise

Progress

댓글 0

'기울기를 추적하지 않는다'고 말하는 세 가지 방법

경계: model.eval()과 no_grad()

Code

External links

Exercise

Progress

경계: `model.eval()`과 `no_grad()`