사용자 정의 Autograd 함수

내장 연산만으로는 부족할 때

모델 코드의 99%는 미분 가능한 내장 연산만으로 충분해. Autograd가 연쇄 법칙에 따라 도함수를 조합하므로 역전파를 직접 작성할 필요가 없지. 나머지 1%에서 torch.autograd.Function이 필요해:

사용자 정의 CUDA 커널처럼 기존의 미분 가능한 연산으로 표현할 수 없는 순전파를 구현할 때.
양자화 인식 학습이나 이산 출력을 내는 모델의 직선 통과 추정기(straight-through estimator)처럼 기울기를 직접 정의해야 할 때.
값을 저장하는 대신 다시 계산해 메모리를 아끼고 싶을 때. 다만 이 용도라면 보통 torch.utils.checkpoint가 더 나아.

계약

사용자 정의 함수에는 두 정적 메서드가 필요해:

forward(ctx, *inputs): 순전파를 계산해. ctx.save_for_backward(...)로 역전파에 필요한 값을 보관해.
backward(ctx, *grad_outputs): 역전파를 계산해. 순전파의 입력마다 기울기 하나를 반환하고, 기울기가 필요 없는 입력에는 None을 반환해.

forward나 backward를 직접 호출하지는 않아. MyFunction.apply(...)를 호출하면 PyTorch가 그래프 노드를 알맞게 구성해.

Code

사용자 정의 ReLU: 이해용, 운영 환경 아님·python

import torch

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)             # stash x for backward
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        grad_in = grad_out.clone()
        grad_in[x < 0] = 0                    # zero gradient where x < 0
        return grad_in

my_relu = MyReLU.apply

x = torch.randn(5, requires_grad=True)
y = my_relu(x)
y.sum().backward()
print(x.grad)   # zeros where x<0, ones where x>=0

Straight-Through Estimator: 기울기 그대로 통과·python

import torch

class STE(torch.autograd.Function):
    """Forward: hard threshold. Backward: pretend it was identity."""
    @staticmethod
    def forward(ctx, x):
        return (x > 0).float()

    @staticmethod
    def backward(ctx, grad_out):
        return grad_out                      # straight through

binarize = STE.apply

# Useful for quantization-aware training:
# the forward step is non-differentiable, but we still need a learning signal
x = torch.randn(4, requires_grad=True)
y = binarize(x)
y.sum().backward()
print(x.grad)   # ones — gradient passed through unchanged

gradcheck로 도함수 검산하기·python

import torch
from torch.autograd import gradcheck

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        grad_in = grad_out.clone()
        grad_in[x < 0] = 0
        return grad_in

# Use float64 for numerical stability when gradchecking
x = torch.randn(8, dtype=torch.float64, requires_grad=True)
print(gradcheck(MyReLU.apply, (x,), eps=1e-6, atol=1e-5))
# True if your analytic backward matches finite-difference numerics

Exercise

사용자 정의 시그모이드 함수의 순전파와 역전파를 모두 구현해 봐. float64 텐서로 torch.autograd.gradcheck를 실행하고, 같은 입력에서 순전파 결과가 torch.sigmoid와 일치하는지도 확인해. 나중에 의심이 생길 때마다 다시 실행할 수 있도록 테스트를 Python 파일로 저장해 둬.

내장 연산만으로는 부족할 때

계약

Code

External links

Exercise

Progress

댓글 0