PyTorch는 어떻게 자동미분을 지원하는가

개요

부스트캠프 AI Tech 7기 과정에서 하루하루를 버텨가던 중 갑자기 포스팅 제목과 같은 질문이 생겼습니다. 그저 럭키 배열에 불과한 줄 알았던 Tensor 객체는 무엇이 그렇게 다르길래 backward() 메소드를 사용하면 알아서 미분을 진행하는 걸까요? 

이게 갑자기 궁금해졌던 저는 공식 문서와 다른 블로그 포스팅을 보면서 스스로가 납득할 수 있는 이유를 찾고자 하였습니다. 이 글을 읽으시는 분들도 같은 생각이라 믿고 제 삽질 경험을 공유해드리겠습니다.

삽질

a = torch.Tensor([1, 2, 3])
a.backward()

일단 모든 Tensor 객체가 backward()를 할 수 있는지 알아보기 위해서 냅다 임의의 텐서를 만들고 바로 backward()를 박아봤습니다. 

그리고 런타임 에러를 받았습니다. 에러 메시지를 읽어보니 아무래도 requires_grad를 True로 바꾸어주어야 하는 모양입니다.

a = torch.Tensor([1, 2, 3])
a.requires_grad = True
a.backward()

a가 미분이 필요하다는 걸 알려주기 위해서 requires_grad를 명시적으로 True로 바꾸어주었습니다.

그리고 다시 한번 더 런타임 에러를 받았습니다. 이번에는 캡처로 같이 보시겠습니다.

어... 일단 해석이 잘 안되는데 에러 메시지를 보아하니 a.numel()이 1이면 저런 문제가 발생하지 않는 것 같습니다. 이게 무슨 함수인지 공식 문서를 찾아보았습니다. 

a.numel() # Returns the total number of elements in the input tensor.

음. 그렇군요. 아무래도 backward()를 사용하고자 하는 함수에는 하나의 원소만 존재해야 하는 모양입니다. 그러고 보면 출력이 하나의 값만 나오는 손실 함수에서 나온 Tensor 객체에 backward() 메소드를 사용했죠. 아무래도 단 한 값에 대한 다른 가중치들의 변화량을 구해야 하는 것이기 때문에 당연하다고 생각할 수 있습니다.

a = torch.Tensor([1, 2, 3])
a.requires_grad = True
b = a.sum()
b.backward()

그럼 간단하게 a의 모든 원소를 다 더한 값 하나만 들고 있는 Tensor 객체 b를 만들어 보았습니다. 그리고 값이 하나뿐인 b에 대해서 backward()를 실행해주면...

오! 성공하였습니다. 근데 jupyter notebook의 출력 셀에서는 아무것도 나오는 게 없습니다? 그래도 메소드가 정상 작동하였으니 변한 내용이 있을 것입니다.

a.grad # tensor([1., 1., 1.])

또 다시 공식 문서를 찾아보니, grad 필드 변수에 접근하여 지금까지 축적된 그래디언트를 확인해볼 수 있다고 합니다. a.grad를 통해 확인하니 실제로 아래 수식처럼 기대한대로 그래디언트를 잘 구했다는 걸 확인할 수 있었습니다.

$$\begin{align}b &= a_1 + a_2 + a_3 \\ \frac{\partial b}{\partial a_1} &= \frac{\partial b}{\partial a_2} = \frac{\partial b}{\partial a_3} = 1  \end{align}$$ 

그러고보니 축적이 된다고 하니까 이 상태에서 바로 또 backward() 함수를 쓰면 그래디언트가 각각 2가 되겠네요?

b.backward()
a.grad # tensor([2., 2., 2.])

네. 실제로 2가 되는 걸 확인할 수 있었습니다. 그래서 그래디언트를 가중치에 반영한 다음에는 opimizer.zero_grad() 함수를 호출하여 이렇게 누적되지 않도록 그래디언트를 계속 지워야 하는 거죠.

그런데 여기서 살짝 인지부조화가 옵니다. 객체 b에 backward() 메소드를 호출했는데 왜 객체 a의 필드 변수에 변화가 생기는 걸까요? 두 객체 사이에는 무언가 연결점이 존재하는 것이 명백해보입니다. 전 그 연결점이 무엇인지 찾고 싶었습니다.

연결점은 어디에

여기서부터는 제 삽질한 내용보다도 깨달은 내용 위주로 글을 적어볼까 합니다. 하나하나 다 적자니 글이 너무 길어지고(그렇게 길게 쓰면 또 제가 너무 힘들고) 그러면 또 너무 글을 읽는 게 지루할테니까요. 그럼 템포를 올려서 가봅시다!

a = torch.Tensor([4])
b = torch.Tensor([3])
a.requires_grad = True
c = a * b
print(c) # tensor([12.], grad_fn=<MulBackward0>)

먼저 requires_grad가 True인 Tensor 객체에 대한 연산을 진행해서 새로운 Tensor 객체를 생성한 뒤, print()문을 통해 내용을 확인해보면 위와 같이 grad_fn에 대한 정보가 새로 추가된 것을 확인해볼 수 있습니다.

f = c.grad_fn

grad_fn 멤버 변수엔 아마 이 공식문서에 나온 클래스의 하위 클래스로 생성된 객체가 저장되는 것으로 보입니다. 그럼 저 객체가 가져야 하는 특성(즉, 저장하고 있어야 하는 필수 정보)이 무엇이 있을지 한번 수식을 통해 알아봅시다.

$$\begin{align}c &= a \times b \\ \frac{\partial c}{\partial a} &= b  \end{align}$$

수식에 의하면 그래디언트를 계산하기 위해서는 b의 값이 필요합니다. 따라서 변수 f에는 Tensor 객체 b에 대한 정보가 담겨 있지 않을까 하는 추론이 가능해집니다.

print(f._saved_self, f._saved_other, sep=", ") # None, tensor([3.])

실제로 dir() 함수를 이용하여 f의 내부 필드 변수 등을 확인해보면, _saved_self나 _saved_other과 같은 무언가 저장하고 있을 것만 같은 변수명을 찾을 수 있습니다. 이를 각각 확인해보면 위의 주석과 같은 값을 갖는다는 확인할 수 있습니다.

_saved_other의 값이 변수 b와 동일한 값을 저장하고 있는 것을 보아하니 우리의 추론이 들어맞은 것을 확인할 수 있습니다. (참고로 코드에선 생략했지만 실제로 두 변수의 id값을 비교하면 동일하다는 것을 알 수 있습니다.)

마무리

연산자마다 객체에 저장되는 값들이 다 다릅니다. 그리고 연산의 층이 깊게 쌓여도 여전히 Tensor 객체들은 기울기가 필요한 최종 Tensor까지 찾아갈 수 있는 방법이 존재합니다.

글의 호흡이 너무 긴 것 같아서 여기서 한 번 끊고 다음 포스팅에서 더 다룰 수 있도록 하겠습니다. 감사합니다.