PyTorch 딥러닝 훈련 메커니즘과 Autograd 이해

딥러닝의 학습 과정은 고차원 공간(파라미터 공간)에서 에너지 포텐셜(Loss Function)이 가장 낮은 지점을 찾아가는 과정과 같습니다. PyTorch는 이 과정을 효율적으로 수행하기 위해 동적 계산 그래프(Dynamic Computational Graph)라는 개념을 사용합니다.

 

MNIST 분류 모델을 예로 들어, 훈련 루프 내부에서 일어나는 일을 단계별로 해부해 보겠습니다.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 1. 데이터셋 준비 (MNIST)
# 텐서 변환 및 정규화
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 학습용 데이터 로드
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 2. 간단한 신경망 모델 정의
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        # 입력(28x28=784) -> 은닉층(128) -> 출력(10: 0~9까지의 숫자)
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 입력 데이터 펼치기 (Flatten)
        x = x.view(-1, 784)
        
        # [Autograd 매커니즘]
        # 이 연산들이 수행될 때 PyTorch는 내부적으로 '계산 그래프'를 동적으로 생성합니다.
        # 각 텐서는 자신이 어떤 연산(Function)을 통해 만들어졌는지 기억합니다 (grad_fn).
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 모델 및 장치 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleNet().to(device)

# 3. 손실 함수와 옵티마이저 정의
criterion = nn.CrossEntropyLoss() # 분류 문제를 위한 손실 함수
optimizer = optim.SGD(model.parameters(), lr=0.01) # 확률적 경사 하강법

# 훈련 함수
def train(model, device, train_loader, optimizer, epoch):
    model.train() # 모델을 훈련 모드로 설정
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        # --- [Step A] 기울기 초기화 ---
        # 이전 배치의 그래디언트가 누적되지 않도록 0으로 초기화
        optimizer.zero_grad()

        # --- [Step B] 순전파 (Forward Pass) ---
        # 입력 데이터를 모델에 통과시켜 예측값(output) 계산
        # 이 과정에서 연산 그래프(Graph)가 구축됨
        output = model(data)

        # --- [Step C] 손실 계산 (Loss Calculation) ---
        # 예측값과 실제값(target) 사이의 오차 계산
        # loss는 그래프의 최종 노드(Leaf Node)가 됨
        loss = criterion(output, target)

        # --- [Step D] 역전파 (Backward Pass - Autograd의 핵심) ---
        # loss.backward() 호출 시 그래프를 역순으로 탐색하며 연쇄 법칙(Chain Rule) 적용
        # 각 파라미터(W, b)들의 .grad 속성에 미분값(기울기)이 저장됨
        loss.backward()

        # --- [Step E] 파라미터 갱신 (Optimizer Step) ---
        # 계산된 기울기(gradient)를 사용하여 모델의 가중치를 수정
        # W_new = W_old - learning_rate * gradient
        optimizer.step()

        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

# 실행
if __name__ == '__main__':
    print(f"Training on {device}")
    for epoch in range(1, 3): # 2 에포크만 실행
        train(model, device, train_loader, optimizer, epoch)

1. 핵심 개념: 텐서(Tensor)와 Autograd

PyTorch의 핵심 데이터 구조인 Tensor는 단순한 다차원 배열(행렬)이 아닙니다. requires_grad=True로 설정된 텐서는 시스템에게 다음과 같이 말합니다.

 

"나를 포함한 모든 연산을 추적해줘. 나중에 미분값이 필요하거든."

 

모델의 가중치(Weight)와 편향(Bias)은 기본적으로 이 옵션이 켜져 있습니다.

2. 훈련 루프의 5단계 메커니즘

Step 1: 순전파 (Forward Pass) - 그래프 구축

코드의 output = model(data) 부분입니다.

• 동작: 입력 데이터 $x$가 행렬 곱셈, 합, 활성화 함수(ReLU) 등을 거쳐 출력 $y$가 됩니다.
• Autograd의 역할:
  • 데이터가 연산을 통과할 때마다 PyTorch는 방향성 비순환 그래프(DAG)를 생성합니다.
  • 입력 텐서는 잎(Leaf) 노드가 되고, 출력 텐서는 루트(Root) 노드가 됩니다.

[심층 분석] .grad_fn: 텐서의 탄생 기록

 

순전파 과정에서 가장 중요한 것은 각 텐서에 grad_fn 속성이 자동으로 부여된다는 점입니다. 물리학적으로 비유하자면, 어떤 입자(텐서)가 관측되었을 때 "이 입자가 어떤 상호작용(덧셈, 곱셈 등)을 통해 생성되었는가?"를 기록한 꼬리표와 같습니다.

1. 연산의 기억 (Provenance):
   
   • 예를 들어 z = x * y라는 연산을 수행하면, z 텐서의 grad_fn은 MulBackward(곱셈의 역전파 함수)라는 객체를 가리킵니다.
   • 만약 z = x + y였다면 AddBackward가 기록됩니다.
2. 연결 리스트 (Linked List):
   
   • 이 grad_fn 객체는 자신을 만든 입력 텐서들의 grad_fn을 다시 참조(next_functions)합니다.
   • 결과적으로 출력층에서 입력층까지 이어지는 거대한 체인(Chain)이 형성됩니다.
3. Leaf Node (잎 노드)의 특이점:
   • 사용자가 직접 생성한 텐서(입력 데이터 x나 가중치 W)는 연산의 결과물이 아니라 시초이므로, grad_fn이 None입니다. 역전파는 여기서 멈추고 계산된 기울기를 .grad에 저장합니다.

Step 2: 손실 계산 (Loss Calculation)

코드의 loss = criterion(output, target) 부분입니다.

• 동작: 예측값과 정답 사이의 차이를 하나의 스칼라 값(Loss)으로 만듭니다.
• 의미: 이 loss 값은 전체 계산 그래프의 최종 종착점입니다. 물리학적으로 보면 현재 시스템의 '불안정성(에러)'을 나타내는 값입니다.

Step 3: 역전파 (Backward Pass) - 자동 미분

코드의 loss.backward() 부분입니다. 가장 마법 같은 부분입니다.

• 동작: loss 텐서에서 시작하여 grad_fn을 따라 그래프를 역순으로 거슬러 올라갑니다.
• 수학적 원리 (연쇄 법칙, Chain Rule):$$\frac{\partial L}{\partial w} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial w}$$
  • PyTorch는 각 grad_fn (예: MulBackward) 내부에 해당 연산의 국소적 미분(Local Gradient) 공식($f'(x)$)을 이미 내장하고 있습니다.
  • 상위 노드에서 내려온 그래디언트에 현재 노드의 국소 그래디언트를 곱하여 하위 노드로 전달합니다.
• 결과: 그래프의 시작점이었던 모델의 파라미터(Weight, Bias)들에 도달하면, 계산된 최종 미분값을 각 파라미터 텐서의 .grad 속성에 저장(누적)합니다.

Step 4: 파라미터 갱신 (Optimizer Step)

코드의 optimizer.step() 부분입니다.

• 동작: 저장된 .grad 값을 사용하여 파라미터 값을 실제로 변경합니다.
• 수식 (SGD의 경우):$$\theta_{new} = \theta_{old} - \eta \cdot \nabla_\theta L$$여기서 $\eta$는 학습률(Learning Rate), $\nabla_\theta L$은 tensor.grad에 저장된 값입니다.
• 물리학적 비유: 언덕(Loss Surface)에서 경사가 가장 가파른 방향(Gradient)의 반대 방향으로 한 발자국 내려가는 것입니다.

Step 5: 그래디언트 초기화 (Zero Grad)

코드의 optimizer.zero_grad() 부분입니다.

• 이유: PyTorch는 기본적으로 .grad에 미분값을 더하는(accumulate) 성질이 있습니다. 이는 RNN 같은 순환 신경망 구현 시 편리하지만, 일반적인 훈련에서는 이전 배치의 미분값이 현재 배치에 영향을 주면 안 되므로 매 루프마다 0으로 초기화해야 합니다.

3. 동적 계산 그래프 시각화

다음 다이어그램은 grad_fn이 실제로 어떻게 연결되어 있는지를 상세히 보여줍니다. 파란색 실선은 데이터의 흐름(Forward)이며, 붉은색 점선은 grad_fn 체인을 타고 흐르는 미분의 흐름(Backward)입니다.

• grad_fn 노드 (보라색): 순전파 때는 연산을 수행하지만, 역전파 때는 미분 계산기 역할을 합니다. 예를 들어 AddmmBackward는 행렬 곱셈과 덧셈에 대한 미분 공식을 알고 있습니다.
• Chain Rule 연결: NllLossBackward $\to$ AddmmBackward $\to$ ReluBackward $\to$ AddmmBackward 순으로 미분값($\frac{\partial L}{\partial y}$)이 전달됩니다.
• .grad 저장 (주황색 점선 박스): 역전파 흐름이 Leaf Node(파라미터)에 도달하면, 계산된 미분값이 해당 파라미터의 .grad 속성에 '주입'됩니다. 연합학습에서는 바로 이 값들을 서버로 보냅니다.