코딩/PyTorch

[PyTorch] Automatic Differentiation with torch.autogrid

guungyul 2025. 1. 10. 22:48

Neural network를 학습할 때 가장 많이 사용되는 알고리즘은 back propagation이다. 이 알고리즘에서는 parameter들이 loss function의 parameter에 대한 gradient에 따라 조절된다.

이런 gradient를 계산하기 위해서 PyTorch에서는 내재된 differentiation 엔진인 torch.autograd를 사용한다.

 

아래 예제에서는 input x, parameters w와 b, 그리고 loss 함수를 사용하는 가장 간단한 one-layer neural network를 보여준다.

import torch

x = torch.ones(5)  # input tensor
y = torch.zeros(3)  # expected output
w = torch.randn(5, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
z = torch.matmul(x, w)+b
loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)

 

Tensors, Functions and Computational graph

위 코드는 다음과 같은 computational graph를 생성한다.

이 network에서 w와 b는 parameter로 최적화를 필요로 한다. 따라서 이 parameters에 대한 loss 함수의 gradient를 계산할 수 있어야 한다. 이를 위해 parameter들의 tensor는 requires_grad 속성을 사용한다.

  • requires_grad는 tensor를 만들때 사용하거나 만든 후 x.requires_grad_(True) 함수를 이용해 정할 수 있음

 

Tensor에 적용해 computational graph를 만드는 함수는 Function class의 객체이다. 이 객체는 forward direction에서 함수를 어떻게 계산해야하는지 알고, backward propagation 단계에서 어떻게 미분값을 구하는지도 알고 있다.

Backward propagation 함수가 참고할 정보는 tensor의 grad_fn 속성에 저장된다.

print(f"Gradient function for z = {z.grad_fn}")
# Gradient function for z = <AddBackward0 object at 0x7f3807d69210>

print(f"Gradient function for loss = {loss.grad_fn}")
# Gradient function for loss = <BinaryCrossEntropyWithLogitsBackward0 object at 0x7f3807d682b0>

 

Computing Gradients

Neural network에 있는 parameter들의 값을 최적화하기 위해서는 각 parameter에 대한 loss function의 미분값을 구해야한다. 이 미분값들을 구하기 위해서 loss.backward() 함수를 부르게 된다.

loss.backward()
print(w.grad)
# tensor([[0.3313, 0.0626, 0.2530],
#         [0.3313, 0.0626, 0.2530],
#         [0.3313, 0.0626, 0.2530],
#         [0.3313, 0.0626, 0.2530],
#         [0.3313, 0.0626, 0.2530]])
        
print(b.grad)
# tensor([0.3313, 0.0626, 0.2530])
  • grad 속성은 computation graph의 leaf node들로부터만 구할 수 있다. 즉, requires_grad 속성이 True로 되어있어야 한다.
  • backward()를 사용한 기울기 계산은 한번만 사용될 수 있다. 만약 같은 graph에 backward()를 여러번 불러야 한다면 backward(retain_graph=True) 속성을 넣어야 한다.

 

Disabling Gradient Tracking

기본적으로 requires_grad=True로 설정된 tensor들은 계산 기록을 저장하고 기울기 계산을 지원한다. 하지만 이런 작업이 필요 없는 경우가 있다. 예를 들어 모델을 몇몇 input data에 대해서 forward computation만 계산하고 싶을 수 있다. 이는 계산 코드를 torch.no_grad() 블록으로 감싸 적용할 수 있다.

z = torch.matmul(x, w)+b
print(z.requires_grad)
# True

with torch.no_grad():
	z = torch.matmul(x, w)+b
print(z.requires_grad)
# False

 

같은 결과를 얻는 다른 방법은 detach() 함수를 사용하는 것이다.

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)
# False

 

Gradient tracking을 비활성화 할 이유는 여러가지가 있을 수 있다.

  • 몇몇 parameter들을 frozen parameter로 만들기 위해
  • Forward 계산을 할 때 계산 속도를 향상

 

More on Computational Graphs

개념적으로 autograd는 데이터 (tensor)들과 실행된 operation들을 Function 객체들로 이루어진 directed acyclic graph (DAG)를 사용해 기록한다. 이 DAG에서 leaves는 input tensor들, roots는 output tensor들로 구성되어 있다.

이 그래프를 root에서 leave로 따라가면서 자동으로 chain rule을 사용한 기울기를 계산할 수 있다.

 

Forward pass에서 autograd는 다음 두 가지를 동시에 진행한다.

  • 결과 tensor를 계산하기 위해 요청된 operation을 실행
  • DAG에 operation의 기울기 함수를 기록

Backward pass는 DAG root에 .backward() 함수가 불리면 실행된다. autograd는 다음 두 가지를 진행한다.

  • 각 .grad_fn에 대해 기울기 계산
  • 대응하는 tensor의 .grad 속성에 값을 저장
  • chain rule을 사용해 leaf tensor들까지 전파

 

Note

  • PyTorch에서 DAG는 동적이다. 즉, 그래프는 처음부터 새로 만들어진다.
  • 각 .backward() 호출 이후 autograd는 새로운 그래프를 채워나간다. 이러한 특징이 사용자가 flow statments를 조정할 수 있게 해준다. 예를 들어 매 iteration마다 shape, size 또는 operation을 바꿀 수 있다.

 

Optional Reading: Tensor Gradients and Jacobian Products

출력 함수가 무작위 tensor인 경우가 있다. 이런 경우에 PyTorch는 실제 기울기가 아닌 Jacobian product를 계산하도록 한다.

Vector 함수 $ \vec{y} = f(\vec{x}) $, $\vec{x} = \langle x_1, \ldots, x_n \rangle$, $ \vec{y} = \langle y_1, \ldots, y_m \rangle$가 있다. 이 함수의 $\vec{x}$에 대한 기울기는 다음 Jacobian matrix에 의해 주어진다.
$$ J = \begin{pmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{pmatrix} $$

 

PyTorch는 Jacobian matrix 자체를 계산하기보다 input vector $ v=(v_1, ..., v_m) $에 대한 Jacobian Product $ v^T * J $를 계산한다. 이는 backward() 함수에 v를 인자로 주어 이루어진다. $v$의 크기는 곱을 구하려고 하는 기존 tensor의 크기와 같아야 한다.

inp = torch.eye(4, 5, requires_grad=True)
out = (inp+1).pow(2).t()
out.backward(torch.ones_like(out), retain_graph=True)
print(f"First call\n{inp.grad}")
# First call
# tensor([[4., 2., 2., 2., 2.],
#         [2., 4., 2., 2., 2.],
#         [2., 2., 4., 2., 2.],
#         [2., 2., 2., 4., 2.]])

out.backward(torch.ones_like(out), retain_graph=True)
print(f"\nSecond call\n{inp.grad}")
# Second call
# tensor([[8., 4., 4., 4., 4.],
#         [4., 8., 4., 4., 4.],
#         [4., 4., 8., 4., 4.],
#         [4., 4., 4., 8., 4.]])
            
inp.grad.zero_()
out.backward(torch.ones_like(out), retain_graph=True)
print(f"\nCall after zeroing gradients\n{inp.grad}")
# Call after zeroing gradients
# tensor([[4., 2., 2., 2., 2.],
#         [2., 4., 2., 2., 2.],
#         [2., 2., 4., 2., 2.],
#         [2., 2., 2., 4., 2.]])

 

backward() 함수가 두번째 호출됐을 때 기울기의 값이 다른 것을 볼 수 있다. 이는 PyTorch가 backward propagation을 할 때 gradient들을 축적하기 때문이다. 예를 들어 계산된 기울기 값들은 모든 leaf node들의 grad 속성에 더해진다.

정확한 기울기를 계산하기 위해서는 grad 속성을 0으로 초기화해야한다. 실제 학습에서는 optimizer가 이 작업을 대신 해준다.

 

Note

  • backward() 함수를 parameter 없이 부르는 것은 backward(torch.tensor(1.0))을 부르는 것과 같다. 이는 scalar-valued 함수의 기울기를 계산하는데에 유용한 방법이다.