리프 대 중간 텐서와 .retain_grad()

기본 설정에서는 리프 텐서만 .grad를 받아

Autograd는 텐서를 두 종류로 뚜렷하게 구분해:

리프 텐서: 직접 만들었거나 nn.Parameter로 등록한 텐서야. 그래프의 입력이지 연산의 출력은 아니야. .grad에는 손실을 이 텐서로 미분한 값이 쌓여.
중간 텐서: 연산의 출력이야. PyTorch는 기본적으로 여기에 기울기를 저장하지 않아. 모든 중간 텐서의 기울기를 보관하면 메모리가 감당하지 못할 만큼 커질 수 있거든.

디버깅할 때 세 번째 계층의 출력처럼 중간 텐서의 기울기를 보고 싶다면 두 가지 방법이 있어:

x.retain_grad(): 역전파 전에 중간 텐서에서 호출하면 PyTorch가 그 텐서의 기울기도 저장해.
x.register_hook(callback): 기울기가 흐를 때 실행할 함수를 등록해. 기울기를 따로 저장하지 않고도 기록하거나 수정할 수 있어.

일반적인 학습 반복문에서는 둘 다 필요하지 않아. '왜 이 계층은 학습되지 않지?'라는 문제가 생겼을 때 진단용으로 꺼내는 도구야.

Code

중간 텐서의 .grad가 None인 이유·python

import torch

a = torch.tensor(2.0, requires_grad=True)   # LEAF
b = a * 3                                    # INTERMEDIATE
c = b + 1                                    # INTERMEDIATE
c.backward()

print(a.is_leaf, a.grad)   # True tensor(3.)
print(b.is_leaf, b.grad)   # False None  (intermediate, .grad not stored)
print(c.is_leaf, c.grad)   # False None

retain_grad: 진단 전용·python

import torch

a = torch.tensor(2.0, requires_grad=True)
b = a * 3
b.retain_grad()                       # tell autograd to keep b's grad
c = b + 1
c.backward()

print(b.grad)                         # tensor(1.) — now available

register_hook: 기울기가 흐를 때 수정하거나 기록하기·python

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2),
)

stats = {}
def make_hook(name):
    def _hook(grad):
        stats[name] = (grad.abs().mean().item(), grad.std().item())
    return _hook

for name, p in model.named_parameters():
    p.register_hook(make_hook(name))

x = torch.randn(16, 4)
y = torch.randn(16, 2)
loss = nn.functional.mse_loss(model(x), y)
loss.backward()

for name, (mean, std) in stats.items():
    print(f"{name}: |grad|.mean={mean:.4f}  std={std:.4f}")

Exercise

3계층 MLP를 만들어 봐. 각 매개변수에 register_hook을 걸어 기울기의 평균과 표준편차를 기록하고, 무작위 데이터 한 배치로 학습한 뒤 통계를 출력해. 그다음 활성화 함수를 ReLU에서 시그모이드로 바꿔 다시 실행해. 입력에 가까운 계층의 기울기 크기가 0에 가까워지는 기울기 소실 현상을 관찰해 봐.

기본 설정에서는 리프 텐서만 .grad를 받아

Code

External links

Exercise

Progress

댓글 0