PyTorch 입문: 텐서·자동미분·학습 루프

지난 글에서 NumPy·Pandas부터 PyTorch·Transformers까지 AI 개발자가 매일 쓰는 Python 라이브러리를 계층별로 훑었다. 이 글은 그중 한 층인 PyTorch를 처음부터 끝까지 한 번 지난다. 텐서와 자동미분, nn.Module로 부품을 만드는 앞쪽 절반과, Dataset·DataLoader와 학습 루프로 그 부품을 엮는 뒤쪽 절반이다.

둘을 한 편에 두는 이유가 있다. 텐서·자동미분·모듈은 부품이고, 학습 루프는 그 부품을 돌리는 조립도다. 부품만 보면 손실 함수와 옵티마이저를 만들어 놓고 한 번도 안 쓴 채 끝나고, 루프만 보면 backward()가 무엇을 채우는지 모른 채 다섯 줄을 외우게 된다. 그래서 앞쪽에서 만든 것이 뒤쪽 루프의 어느 줄에 들어가는지를 계속 짚으며 간다. 다 읽고 나면 optimizer.zero_grad()를 왜 매 배치 부르는지, model.eval()과 torch.no_grad()가 왜 둘 다 필요한지가 규칙이 아니라 이유로 남는다.

동적 그래프와 텐서

동적 계산 그래프

PyTorch는 2016년 9월 Facebook AI Research(FAIR, 지금의 Meta AI)가 첫 공개 릴리스를 낸 딥러닝 프레임워크다. 가장 큰 설계 결정은 동적 계산 그래프(define-by-run)다. 계산 그래프는 어떤 텐서에 어떤 연산을 걸어 어떤 텐서가 나왔는지를 노드와 간선으로 적어 둔 기록이고, 역전파는 이 기록을 거꾸로 따라간다. 초기 TensorFlow는 그래프를 먼저 통째로 정의하고 나중에 세션에서 실행하는 정적 계산 그래프(define-and-run)를 썼다. 그래프를 미리 알면 최적화하기 좋지만, 파이썬 코드가 실제 계산이 아니라 계산의 설계도를 짜는 셈이라 중간 값을 찍어 볼 수도 보통의 조건문을 걸 수도 없었다.

PyTorch는 반대로 파이썬 문장이 실행되는 순간 연산이 일어나고, 그래프는 그 실행을 따라가며 만들어진다. 그래서 그래프가 순전파마다 새로 생기고 backward()가 끝나면 버려진다. 배치마다 다른 모양의 그래프가 있어도 상관없다는 뜻이고, 길이가 제각각인 시퀀스나 입력에 따라 갈라지는 분기를 다루기가 자연스럽다.

동적 그래프의 실용적 이득은 두 가지다. 첫째, 파이썬의 if와 for를 모델 코드에 그대로 쓴다. 토큰 길이만큼 도는 for, 어떤 조건에서만 타는 분기, 재귀 구조가 전부 보통 파이썬으로 적힌다. 정적 그래프에서는 이런 것을 프레임워크 전용 제어 연산으로 바꿔 적어야 했다. 둘째, 디버깅이 보통 파이썬 디버깅과 같다. forward 중간에 print(x.shape)를 넣으면 그 자리의 모양이 찍히고, 예외가 나면 스택 트레이스가 그 줄을 가리킨다. 모양이 안 맞아 터지는 자리를 찾는 일이 딥러닝 개발 시간의 큰 몫인데, 이것이 곧바로 되는 것과 그래프를 다시 짜서 확인해야 하는 것은 하루의 길이가 다르다.

그 편의가 연구 쪽을 먼저 끌어왔다. 논문 구현 코드를 모아 세는 집계에서 PyTorch 비율이 오랫동안 다수를 차지했고, 새 구조를 빠르게 시험하는 자리에서는 사실상 표준이 됐다. 만든 곳은 한 회사지만 관리는 이미 넘어가서, 지금은 리눅스 재단이 호스팅하는 PyTorch Foundation이 벤더 중립으로 이 프로젝트를 맡는다. 이 글에서 다루는 것은 그 표준의 가장 아래층이다 — 텐서 하나에서 시작해 루프 한 바퀴까지.

NumPy 배열과의 차이

PyTorch의 기본 자료형은 torch.Tensor다. 텐서는 숫자를 여러 차원으로 늘어놓은 배열이고, 0차원이면 스칼라, 1차원이면 벡터, 2차원이면 행렬이다. 차원이 더 올라가도 이름은 텐서 하나로 통일된다. 이미지 배치는 (배치, 채널, 높이, 너비)의 4차원, 비디오는 여기에 시간축이 하나 더 붙은 5차원이다. shape가 그 차원의 크기를 튜플로 알려 주고, 모양을 읽는 습관이 이 프레임워크를 다루는 기본기의 절반이다.

NumPy의 ndarray를 써 봤다면 텐서는 거의 같은 물건이다. 인덱싱, 브로드캐스팅, 원소별 연산이 같은 규칙으로 돌아간다. 다른 것은 둘뿐인데 그 둘이 딥러닝을 가능하게 한다. 하나는 GPU 연산이다. 같은 코드가 device만 바꾸면 GPU에서 돈다. 다른 하나는 자동미분이다. 텐서에 걸린 연산이 기록되고, 그 기록으로 기울기를 자동으로 계산한다. NumPy 배열은 값만 들고 있지만 텐서는 값과 함께 「이 값이 어디서 왔는가」를 들고 있다.

생성·속성·장치 이동

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])   # 값을 직접 지정
z = torch.zeros(3, 4)                          # 0으로 채운 (3, 4)
r = torch.randn(2, 3)                          # 표준정규분포 난수
e = torch.ones_like(x)                         # x와 같은 모양, 전부 1

print(x.shape, x.dtype, x.device)   # torch.Size([2, 2]) torch.float32 cpu

device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)                    # 옮긴 텐서가 새로 돌아온다 (x.cuda()도 같다)

텐서에는 세 속성이 늘 따라다닌다. shape는 모양, dtype는 원소의 자료형, device는 값이 놓인 장치다. 셋 중 오류를 가장 자주 내는 것은 device다. 연산에 참여하는 텐서는 전부 같은 장치에 있어야 하고, 하나라도 CPU에 남아 있으면 「모든 텐서가 같은 장치에 있어야 한다」는 런타임 오류가 난다. 모델을 GPU로 보냈는데 배치를 안 보냈거나, 배치는 보냈는데 손으로 만든 마스크 텐서 하나를 빠뜨리는 식이다. 그래서 학습 루프의 첫 줄은 언제나 x, y = x.to(device), y.to(device)이고, 이 줄은 뒤에서 그대로 다시 나온다.

.to()가 새 텐서를 돌려준다는 것도 처음에 한 번은 밟는 자리다. x.to(device)라고만 쓰고 결과를 받지 않으면 x는 그대로 CPU에 있다. 반면 model.to(device)는 모듈 안의 파라미터를 제자리에서 옮기므로 결과를 받지 않아도 된다. 텐서는 값이고 모듈은 그릇이라 동작이 다르다.

dtype는 기본이 float32다. 정수 리스트로 만든 텐서는 int64가 되는데, 정수 텐서는 기울기를 가질 수 없어서 requires_grad를 켜면 그 자리에서 거부된다. 실수 연산에 정수 텐서가 섞여 들어가 자료형 오류가 나면 만들 때 1.0처럼 소수점을 붙였는지부터 본다.

모양 변경 연산

a = torch.randn(3, 4)
b = torch.randn(4, 5)

c = a @ b             # 행렬 곱 → (3, 5)
d = a.T               # 전치 → (4, 3)
e = a.unsqueeze(0)    # (1, 3, 4) — 배치 차원을 앞에 끼운다
f = e.squeeze(0)      # (3, 4)    — 크기 1인 차원을 뺀다
g = a.reshape(2, 6)   # 메모리가 연속이 아니어도 된다
h = a.view(2, 6)      # 연속 메모리일 때만, 복사 없이

산술은 NumPy와 같지만 딥러닝에서 유난히 자주 쓰는 넷은 따로 익혀 둔다. @는 행렬 곱이고, 앞 텐서의 마지막 차원과 뒤 텐서의 첫 차원이 같아야 한다. (3, 4)와 (4, 5)를 곱해 (3, 5)가 나오는 식이다. .T는 전치다. unsqueeze와 squeeze는 크기 1인 차원을 끼우고 빼는 짝인데, 모델이 배치를 전제로 짜여 있어 샘플 하나를 넣을 때 (3, 4)를 (1, 3, 4)로 만들어야 하는 자리가 반드시 온다. 배치는 한 번에 함께 처리하는 샘플 묶음이고, 배치 차원은 그 묶음을 세려고 맨 앞에 두는 차원이다. 모델 코드는 이 차원이 있다고 가정한다.

reshape와 view는 같은 일을 하는 것처럼 보이고 실제로 결과도 같다. 차이는 메모리다. view는 원래 텐서의 메모리를 그대로 다른 모양으로 읽는 창이라 복사가 없고 빠르지만, 원소가 메모리에 차례대로 놓여 있을 때만 된다. 전치를 한 번 거친 텐서는 논리적 순서와 메모리 순서가 어긋나서 view가 거부한다. reshape는 가능하면 view처럼 동작하고 안 되면 복사해서라도 만들어 준다. 습관으로는 reshape를 쓰고, 성능을 재는 자리에서 복사가 일어나는지 볼 때만 view로 바꿔 확인하는 편이 안전하다.

autograd

연산 기록

신경망 학습은 역전파(backpropagation)로 돌아간다. 손실을 각 파라미터로 미분한 값, 곧 기울기(gradient)를 구해 파라미터를 그 반대 방향으로 조금 옮기는 일의 반복이다. 손으로 미분식을 적어 넣던 시절도 있었지만, 층이 수십 개인 모델에서 그 일을 사람이 하는 것은 불가능에 가깝다. autograd는 PyTorch의 자동미분 엔진으로, 순전파 동안 연산을 기록해 두었다가 역전파 때 그 기록을 거슬러 기울기를 계산한다.

기록의 스위치가 requires_grad다. 이 값이 True인 텐서가 연산에 참여하면 그 연산과 결과가 계산 그래프에 올라가고, 결과 텐서도 requires_grad=True가 되어 다음 연산도 이어서 기록된다. 그래프의 끝에서 backward()를 부르면 엔진이 간선을 거꾸로 따라가며 기울기를 채운다. 모델의 파라미터는 nn.Module이 만들 때 이미 requires_grad=True라 우리가 켤 일은 드물다. 손으로 켜는 자리는 입력 자체를 최적화할 때 정도다 — 적대적 예제를 만들거나 어떤 입력이 뉴런을 가장 크게 켜는지 찾을 때 그렇다.

backward와 연쇄 법칙

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x    # y = x² + 3x
y.backward()          # dy/dx = 2x + 3
print(x.grad)         # tensor([7.])

$y = x^2 + 3x$ 를 $x = 2$ 에서 미분하면 $2x + 3 = 7$ 이고, x.grad에 정확히 그 값이 들어온다. 이 작은 예에서 세 가지를 읽을 수 있다. 첫째, 우리는 미분식을 어디에도 적지 않았다. 제곱과 곱과 덧셈 세 연산이 기록됐고, 엔진은 각 연산의 도함수를 알고 있다. 둘째, backward()는 스칼라에서 부른다. 손실이 언제나 숫자 하나로 줄어드는 이유가 여기 있다 — 벡터에서 부르려면 어느 방향으로 미분할지를 따로 넘겨야 한다. 셋째, 기울기는 그래프의 잎, 곧 requires_grad를 켜고 직접 만든 텐서에만 남는다. 중간 결과인 y에는 .grad가 없다.

연쇄 법칙(chain rule)이 이 엔진의 수학 전부다. 합성함수의 미분은 각 단계 도함수의 곱이라는 규칙이고, 그래프를 거꾸로 걸으며 도함수를 곱해 나가는 것이 그 규칙의 기계적 실행이다. 층이 수십 개여도 각 층은 자기 도함수만 알면 되고, 곱하는 일은 엔진이 한다. 그래서 새 연산이나 새 층을 만들 때도 순전파와 그 연산 하나의 도함수만 정의하면 나머지는 저절로 이어진다.

기울기 누적

x.grad는 backward()가 덮어쓰지 않고 더한다. 위 예에서 y를 다시 계산해 backward()를 한 번 더 부르면 14가 되고, 또 부르면 21이다. 처음 보면 버그 같지만 의도된 설계다. 메모리가 모자라 배치를 둘로 쪼개 순전파와 역전파를 따로 돌리고 한 번만 파라미터를 갱신하는 기울기 누적(gradient accumulation)이 이 성질 위에 서 있고, 여러 손실을 따로 역전파해 합치는 것도 마찬가지다.

대가는 우리가 매번 비워야 한다는 것이다. 텐서 하나면 x.grad.zero_()이고, 모델 전체면 optimizer.zero_grad()다. 학습 루프에서 이 호출을 매 배치 넣는 이유가 바로 이것이다 — 안 비우면 이번 배치의 기울기에 지난 배치의 기울기가 더해지고, 그 위에 그 앞 배치가 더해져서, 몇 배치 뒤에는 파라미터가 엉뚱한 방향으로 튄다. 오류는 나지 않는다. 손실이 내려가다 말거나 발산할 뿐이라, 이 줄 하나가 빠진 것을 찾는 데 하루가 간다.

no_grad와 추론

기록은 공짜가 아니다. 그래프를 유지하려면 역전파에 쓸 중간 결과를 전부 메모리에 잡아 두어야 하고, 이것이 학습 중 GPU 메모리의 큰 몫이다. 기울기가 필요 없는 자리, 곧 검증과 추론에서는 기록을 끄는 것이 맞다. torch.no_grad() 컨텍스트 안에서는 requires_grad가 켜진 텐서로 연산해도 그래프에 오르지 않는다.

with torch.no_grad():
    pred = model(inputs)

효과가 둘이다. 중간 결과를 안 잡아 두니 메모리가 줄고, 같은 이유로 검증 배치를 학습 배치보다 크게 잡을 수 있다. 그리고 기록 자체의 오버헤드가 없어 조금 빨라진다. 뒤에서 검증 함수를 볼 때 model.eval()과 짝으로 나오는데, 둘이 하는 일이 다르다는 점을 그때 다시 짚는다.

nn.Module

생성자와 forward의 분담

torch.nn.Module은 모든 PyTorch 모델의 부모 클래스다. 층 하나도 모듈이고 층을 묶은 모델도 모듈이라, 모듈 안에 모듈이 들어가는 나무 구조가 된다. 모듈을 쓸 때 우리가 적는 것은 두 메서드뿐이다. 생성자 __init__에서 무엇을 갖고 있는가를 적고, forward에서 데이터가 어떤 순서로 지나가는가를 적는다.

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x):
        return self.layers(x)

model = MLP(784, 256, 10)

두 메서드를 가르는 데는 이유가 있다. __init__에서 nn.Linear를 self. 뒤에 붙여 두면 부모 클래스가 그것을 등록한다. 등록된 모듈의 파라미터가 model.parameters()에 잡히고, model.to(device)에 함께 옮겨지며, state_dict()에 함께 저장된다. 이 등록이 모듈 체계의 핵심이다. 층을 파이썬 리스트에 담아 두면 등록이 안 되어 옵티마이저가 그 파라미터를 못 보고, 학습이 돌아가는 것처럼 보이면서 그 층만 갱신되지 않는다. 층 여러 개를 목록으로 들고 있어야 하면 nn.ModuleList를 쓴다.

forward는 직접 부르지 않는다. model(x)라고 쓰면 부모 클래스의 __call__이 forward를 대신 불러 주는데, 그 앞뒤에 훅이 걸려 있어서 forward를 직접 부르면 그 훅이 빠진다. 규칙으로 외우면 「정의는 forward에, 호출은 model(x)로」다.

Sequential과 분기 구조

위 코드에서 forward가 한 줄로 끝나는 것은 nn.Sequential이 층을 순서대로 이어 주기 때문이다. Sequential은 넣은 순서대로 앞 층의 출력을 다음 층의 입력으로 넘기는 컨테이너다. 입력이 하나로 들어와 하나로 나가는 직선 구조면 이것으로 충분하고, 층 순서가 코드에 그대로 보여서 읽기도 쉽다.

직선이 아닐 때 forward를 손으로 쓴다. 입력을 두 갈래로 보냈다가 합치는 분기, 몇 층을 건너뛰어 더하는 잔차 연결(residual connection), 입력 둘을 받는 모델, 조건에 따라 다른 경로를 타는 구조가 그렇다. 순서대로 흘리는 것만으로는 표현이 안 되는 자리다. 이때도 __init__은 그대로다 — 층은 여전히 거기서 만들고, forward에서 그 층들을 어떤 순서와 조합으로 부를지만 달라진다. 그림의 MLP처럼 fc1·relu·fc2를 따로 두고 forward에서 차례로 부르는 것과 Sequential에 넣는 것은 완전히 같은 모델이고, 뒤에 분기를 넣을 계획이면 앞의 방식이 손이 덜 간다.

forward 안에서 활성화 함수는 두 가지로 쓸 수 있다. nn.ReLU()처럼 모듈로 만들어 두거나 torch.relu(x)처럼 함수로 그 자리에서 부르거나. 파라미터가 없는 연산이라 어느 쪽이든 결과는 같다. 다만 Sequential에 넣으려면 모듈이어야 하고, 학습·추론에서 동작이 달라지는 것은 모듈로 두어야 model.eval()이 그것을 찾아 끌 수 있다. 드롭아웃이 그런 층이다 — 학습 중 뉴런 일부를 무작위로 꺼서 특정 뉴런에 기대지 못하게 하는 정규화이고, 추론 때는 전부 켜야 하므로 모드에 따라 동작이 갈린다.

파라미터와 state_dict

total = sum(p.numel() for p in model.parameters())
print(f"파라미터 수: {total:,}")   # 203,530

torch.save(model.state_dict(), "model.pt")
model.load_state_dict(torch.load("model.pt"))

parameters()는 등록된 모든 파라미터 텐서를 차례로 내주고, numel()은 텐서 하나의 원소 수다. 둘을 합치면 모델 크기가 나온다. MLP(784, 256, 10)이면 첫 층이 가중치 $784 \times 256 = 200{,}704$ 개에 편향 256개, 둘째 층이 가중치 $256 \times 10 = 2{,}560$ 개에 편향 10개라 모두 203,530개다. 새 구조를 짤 때 이 수를 한 번 찍어 보는 습관이 값이 있다. 층 하나를 잘못 이어 파라미터가 열 배로 뛰거나 등록이 안 되어 0으로 나오는 것을 이 한 줄이 바로 보여 준다.

state_dict는 모듈이 가진 파라미터와 버퍼를 이름 → 텐서의 사전으로 꺼낸 것이다. 저장하는 것은 모델 객체가 아니라 이 사전이다. 모델 객체를 통째로 피클하면 클래스 정의 경로까지 같이 묶여서 코드를 조금만 옮겨도 불러오기가 깨지지만, 사전은 값뿐이라 같은 구조의 모델을 새로 만들고 load_state_dict로 채우면 된다. 사전의 키가 layers.0.weight처럼 모듈 경로를 따르므로, 구조를 바꾸면 키가 안 맞아 그 자리에서 오류가 난다. 조용히 틀리는 것보다 낫다.

버퍼가 여기 함께 들어간다는 점을 기억해 둔다. 버퍼는 학습되지 않지만 모델 상태의 일부인 텐서로, 층의 입력을 배치 통계로 정규화해 학습을 안정시키는 배치 정규화 층이 들고 다니는 이동 평균과 분산이 대표다. 파라미터만 저장하면 그 통계가 빠져 추론 결과가 달라지는데, state_dict는 둘을 같이 담으므로 이 경로를 따르는 한 걱정할 일이 없다.

내장 층

층	하는 일
nn.Linear(in, out)	완전 연결층, $Wx + b$
nn.Conv2d(in, out, k)	2차원 합성곱
nn.BatchNorm1d/2d	배치 정규화
nn.Dropout(p)	드롭아웃
nn.Embedding(V, d)	정수 인덱스 → 벡터
nn.LSTM(in, h)	LSTM 층
nn.MultiheadAttention	멀티헤드 어텐션

torch.nn에는 논문에 나오는 층 대부분이 이미 있다. 표의 일곱은 그중 어느 분야를 가든 만나는 것들이다. 완전 연결층과 합성곱과 순환층은 각각 표 형태·이미지·시퀀스 데이터의 기본 부품이고, 임베딩은 단어 번호 같은 정수를 학습 가능한 벡터로 바꾸는 조회표이며, 배치 정규화와 드롭아웃은 학습을 안정시키고 과적합(모델이 학습 데이터에만 맞아 새 데이터에서 무너지는 것)을 막는 보조 층이다. 어텐션은 Transformer의 심장이다. 어느 것이든 __init__에서 만들어 forward에서 부르는 방식이 같아서, 새 층을 익힐 때 볼 것은 생성 인자와 입력·출력 모양 둘뿐이다.

손실 함수와 옵티마이저는 모델 밖에 따로 둔다. 손실 함수는 예측과 정답 사이의 거리를 숫자 하나로 만드는 함수이고, 옵티마이저는 그 숫자의 기울기로 파라미터를 옮기는 규칙이다. 둘 다 모델과 독립이라 같은 모델에 다른 손실을 걸거나 옵티마이저만 바꿔 실험한다.

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

옵티마이저는 만들 때 model.parameters()를 받는다. 이 순간 옵티마이저가 갱신할 텐서 목록이 정해지므로, 그 뒤에 모델에 층을 더하면 새 층은 갱신 대상이 아니다. 그리고 nn.CrossEntropyLoss는 소프트맥스를 안에 품고 있어 모델의 마지막 층은 소프트맥스 없이 점수를 그대로 내보내야 한다. 이 점수를 로짓(logit)이라 부른다. 마지막에 소프트맥스를 한 번 더 걸면 오류 없이 학습이 느려지기만 하는데, 분류 모델을 처음 짤 때 흔히 밟는 자리다.

여기까지가 부품이다. 모델이 있고 손실 함수와 옵티마이저가 있는데, 아직 아무것도 학습하지 않았다. 데이터를 배치로 흘려 넣는 장치와 그것을 돌리는 루프가 남았다.

Dataset과 DataLoader

Dataset

PyTorch에서 데이터를 다루는 표준은 Dataset이 샘플 하나를, DataLoader가 배치를 맡는 2단 구조다. Dataset은 「샘플이 몇 개인가」와 「i번째 샘플이 무엇인가」 두 질문에 답하는 객체이고, 그래서 구현할 메서드도 __len__과 __getitem__ 둘뿐이다. 파일에서 읽든 메모리에 들고 있든 API를 호출하든 그 두 메서드 뒤로 숨는다.

from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.encodings = tokenizer(texts, truncation=True, padding=True,
                                   max_length=max_len, return_tensors="pt")
        self.labels = torch.tensor(labels)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return {k: v[idx] for k, v in self.encodings.items()}, self.labels[idx]

이 예는 텍스트 분류용이다. 토크나이저는 문장을 모델이 받는 정수 번호 열로 바꾸는 도구인데, 생성 시점에 문장 전체를 한 번에 바꿔 두고 __getitem__은 그 결과에서 idx번째 행을 잘라 라벨과 함께 돌려준다. truncation은 최대 길이를 넘는 문장을 자르고 padding은 짧은 문장을 채워 모든 행을 같은 길이로 맞추는 옵션이고, 그래서 배치로 묶을 때 모양이 맞는다.

미리 바꿔 두는 방식과 __getitem__에서 그때그때 바꾸는 방식은 서로 다른 것을 치른다. 미리 바꾸면 학습 중 CPU가 놀지만 전체 데이터가 텐서로 메모리에 올라가야 하고, 그때그때 바꾸면 메모리는 가볍지만 에폭(학습 데이터 전체를 한 번 다 지나는 단위)마다 같은 변환을 되풀이한다. 데이터가 메모리에 다 들어가면 앞의 방식이 단순하고, 이미지처럼 크거나 무작위 증강을 걸어야 해서 매번 달라야 하면 뒤의 방식이다. 어느 쪽이든 Dataset의 바깥 모양은 같아서 DataLoader는 차이를 모른다.

배치·셔플·프리페치

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=4)
val_loader   = DataLoader(val_ds,   batch_size=64, shuffle=False, num_workers=4)

DataLoader는 Dataset에서 샘플을 하나씩 꺼내 배치로 쌓고, 순서를 섞고, 별도 프로세스에서 미리 준비해 두는 장치다. 세 인자가 그 세 일을 정한다. batch_size는 한 번에 묶을 샘플 수다. shuffle은 에폭마다 순서를 섞을지인데, 학습 데이터는 반드시 섞고 검증 데이터는 섞지 않는다. 섞지 않으면 같은 클래스가 파일 순서대로 뭉쳐 들어와 배치마다 기울기가 한쪽으로 쏠리고, 검증은 순서가 결과에 영향을 주지 않으니 섞을 이유가 없으면서 섞으면 재현이 어려워진다.

검증 배치를 학습 배치의 두 배로 잡은 것은 앞에서 본 no_grad의 결과다. 검증에서는 역전파용 중간 결과를 잡아 두지 않으니 같은 메모리에 두 배를 넣을 수 있고, 배치가 크면 배치 수가 줄어 검증이 빨리 끝난다. 학습 배치 크기는 메모리 한도와 학습 안정성이 함께 정하는 값이라 이렇게 자유롭지 않다.

배치를 쌓는 일은 기본 collate 함수가 맡는다 — 샘플 여러 개를 받아 배치 하나로 묶는 함수다. __getitem__이 텐서를 주면 새 배치 차원으로 쌓고, 사전을 주면 키마다 따로 쌓아 사전으로 돌려준다. 위의 TextDataset이 사전과 라벨의 튜플을 주므로 루프에서 for x, y in loader로 받으면 x가 배치 차원이 붙은 사전이 되고, 장치로 옮길 때도 {k: v.to(device) for k, v in x.items()}처럼 값마다 옮긴다. 길이가 제각각인 시퀀스를 패딩 없이 다루려면 이 함수를 직접 넘겨야 하는데, 여기서는 Dataset이 이미 길이를 맞춰 두어 기본으로 충분하다.

num_workers와 GPU 대기

num_workers=4는 데이터를 준비하는 프로세스를 넷 띄우라는 뜻이다. GPU가 한 배치를 계산하는 동안 CPU가 다음 배치를 읽고 변환해 두면 GPU가 데이터를 기다리지 않는다. 이 값이 0이면 메인 프로세스가 배치를 만들고 나서야 GPU를 부르므로 두 장치가 번갈아 놀게 된다. 이미지 디코딩이나 증강처럼 샘플당 CPU 일이 무거운 데이터에서는 이 인자 하나가 학습 시간을 몇 배로 가른다.

병목이 어디 있는지는 GPU 사용률로 본다. 학습 중 사용률이 오르내리며 자주 바닥을 치면 GPU가 데이터를 기다리는 것이고, 이때 모델을 손보는 것은 헛일이다. 워커 수를 늘리거나 Dataset의 변환을 가볍게 하는 것이 맞다. 반대로 사용률이 계속 높으면 데이터 쪽은 문제가 아니다. 워커를 CPU 코어 수보다 많이 띄워도 이득이 없고 메모리만 든다는 것, 그리고 워커마다 Dataset 객체가 복제되므로 __init__에서 큰 것을 들고 있으면 그 수만큼 메모리를 쓴다는 것을 함께 기억해 둔다.

학습 루프

한 배치의 다섯 줄

부품이 다 모였다. 모델, 손실 함수, 옵티마이저, 그리고 배치를 내주는 로더. 학습 루프는 이 넷을 정해진 순서로 부르는 것이고, 배치 하나에 하는 일은 다섯 줄이다.

model     = MyModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0.0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()                              # ① 기울기 비우기
        pred = model(x)                                    # ② 순전파
        loss = criterion(pred, y)                          # ③ 손실
        loss.backward()                                    # ④ 역전파
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)  #    (클리핑)
        optimizer.step()                                   # ⑤ 파라미터 갱신
        total_loss += loss.item()
    return total_loss / len(loader)

다섯 줄을 앞 절들과 이어 읽는다. ①은 autograd 절에서 본 누적 문제를 막는 줄이다. ②는 nn.Module 절의 forward가 __call__을 거쳐 불리는 자리이고, 여기서 계산 그래프가 새로 만들어진다. ③은 예측과 정답의 거리를 스칼라 하나로 줄인다 — backward()를 스칼라에서만 부를 수 있다고 한 그 스칼라다. ④가 그래프를 거꾸로 걸으며 모든 파라미터의 .grad를 채우고, ⑤에서 옵티마이저가 그 .grad를 읽어 파라미터를 옮긴다. 옵티마이저가 만들 때 받아 둔 model.parameters()가 여기서 쓰인다. 모델이 무엇이든 이 루프는 같다 — MyModel이 앞의 MLP여도, 순환층이나 Transformer여도 다섯 줄은 한 글자도 안 바뀐다.

옵티마이저는 앞 절의 Adam이 아니라 AdamW로 바꿨다. 가중치 감쇠(weight decay)는 파라미터가 커지는 것을 매 스텝 조금씩 눌러 과적합을 막는 정규화이고, AdamW는 그 감쇠를 Adam의 적응적 학습률과 분리해 적용하는 변형이다. Transformer 계열에서 기본값처럼 쓰이는 조합이라 예도 그쪽으로 맞췄다. 학습률 3e-4는 근거가 있는 상수는 아니고 Adam 계열에서 첫 시도로 자주 쓰는 값이다.

model.train()을 루프 앞에 부르는 것을 빠뜨리기 쉽다. 검증 함수가 model.eval()로 바꿔 놓은 것을 되돌리는 줄이라, 첫 에폭에는 없어도 티가 안 나다가 둘째 에폭부터 드롭아웃이 꺼진 채 학습한다.

zero_grad의 자리

①이 루프의 시작에 있는 것은 우연이 아니다. 기울기는 backward()가 채우고 step()이 읽으므로, 비우는 일은 이번 배치의 backward() 앞이기만 하면 된다. 그림처럼 손실 계산 뒤에 두어도 맞다. 하지만 루프의 끝에 두면 사정이 달라진다. 이번 배치의 step() 뒤에 비우고 다음 배치가 그 빈 상태를 이어받는 구조인데, 이 구조는 루프 바깥에서 기울기가 생기지 않는다는 가정에 기댄다. 루프에 들어오기 전에 디버깅으로 backward()를 한 번 불렀거나, 체크포인트에서 재개했거나, 검증 코드가 실수로 no_grad 밖에서 돌았다면 첫 배치가 그 기울기를 떠안는다. 시작에 두면 그런 가정이 필요 없다.

zero_grad()가 하는 일은 등록된 파라미터의 .grad를 비우는 것이 전부다. 여러 옵티마이저를 쓰거나 파라미터 일부만 학습하는 자리에서는 옵티마이저가 모르는 파라미터의 .grad가 남으므로 model.zero_grad()로 모델 쪽에서 비우는 편이 확실하다.

기울기 클리핑

clip_grad_norm_은 다섯 단계 사이에 끼운 안전장치다. 기울기 클리핑(gradient clipping)은 모든 파라미터의 기울기를 하나의 벡터로 보고 그 크기(노름)가 한계를 넘으면 한계에 맞게 줄이는 조작이다. 위 코드의 1.0이 그 한계다. 방향은 그대로 두고 크기만 자르므로 「어느 쪽으로 갈지」는 유지하고 「얼마나 갈지」만 막는다.

자리가 정해져 있다. backward()가 기울기를 채운 뒤, step()이 그것을 읽기 전이다. 순서를 바꾸면 자르지 않은 기울기로 갱신하고 나서 자르는 셈이라 효과가 없다. 필요한 모델도 정해져 있다. 순환 신경망과 Transformer처럼 같은 가중치를 시퀀스 길이만큼 되풀이해 곱하는 구조는 기울기가 지수적으로 커지는 기울기 폭발이 쉽게 일어나고, 한 배치의 폭발이 파라미터를 망가뜨리면 그 뒤로는 손실이 nan으로 굳는다. 클리핑은 그 한 배치를 막는다. MLP 같은 얕은 모델에서는 없어도 되지만 넣어서 손해 볼 것도 없어 습관으로 두는 사람이 많다.

item()과 손실 누적

마지막 줄 total_loss += loss.item()이 사소해 보이지만 .item()이 빠지면 학습이 에폭 중간에 메모리 부족으로 죽는다. loss는 계산 그래프에 매달린 텐서라, 파이썬 실수 변수에 loss를 그대로 더하면 total_loss가 텐서가 되고 그래프에 매달린 채 남는다. 배치마다 그래프 하나씩이 해제되지 못하고 쌓이는 것이다. .item()은 원소 하나짜리 텐서에서 파이썬 숫자만 꺼내므로 그래프와의 연결이 끊긴다. 로깅·기록·비교에 쓰는 값은 전부 .item()이나 .detach()를 거쳐 그래프에서 떼어 낸다.

같은 이유로 .item()을 남발하지도 않는다. GPU에서 값을 꺼내는 것은 GPU가 그 값을 다 계산할 때까지 CPU를 세우는 동기화 지점이라, 배치마다 여러 번 부르면 그만큼 느려진다. 배치당 한 번, 에폭 평균을 내는 데 쓰는 정도가 적당하다. 반환하는 값은 배치 손실의 평균이고 len(loader)는 배치 수이므로, 마지막 배치가 작아도 대략 맞는다.

검증과 체크포인트

eval 모드와 no_grad의 분업

def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            pred = model(x)
            total_loss += criterion(pred, y).item()
            correct    += (pred.argmax(1) == y).sum().item()
            total      += y.size(0)
    return total_loss / len(loader), correct / total

검증 함수의 첫 두 줄이 model.eval()과 torch.no_grad()인데, 둘은 서로 다른 것을 끈다. model.eval()은 모듈의 동작 모드를 바꾼다. 드롭아웃은 학습 때 뉴런을 무작위로 끄고 추론 때는 전부 켜야 하며, 배치 정규화는 학습 때 그 배치의 통계로 정규화하고 추론 때는 학습 중 모아 둔 이동 평균과 분산을 써야 한다. 이 전환을 eval()이 한다. 안 하면 추론 결과가 배치 구성에 따라 달라지고 드롭아웃이 켜진 채 예측이 나와 정확도가 낮게 잡힌다.

torch.no_grad()는 autograd를 끈다. 모듈은 아무것도 모르고 계산은 그대로 하되 기록만 안 한다. 그래서 하나만 쓰면 반쪽이다. eval()만 쓰면 결과는 맞지만 메모리를 학습 때만큼 쓰고, no_grad()만 쓰면 메모리는 아끼지만 드롭아웃이 켜진 채 검증한다. 학습이 잘 되는데 검증 정확도가 이상하게 낮으면 이 두 줄부터 본다.

정확도는 argmax로 잰다. 모델이 클래스마다 점수를 내므로 가장 큰 점수의 인덱스가 예측 클래스이고, 정답과 같은 개수를 세어 전체로 나눈다. 손실은 배치 평균의 평균이고 정확도는 샘플 단위의 비율이라 분모가 다르다는 것을 눈여겨본다.

조기 종료와 최선의 모델

best_val_loss = float("inf")
patience, no_improve = 5, 0

for epoch in range(1, num_epochs + 1):
    train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    scheduler.step(val_loss)                          # ReduceLROnPlateau
    print(f"[{epoch:03d}] train={train_loss:.4f} val={val_loss:.4f} acc={val_acc:.4f}")

    if val_loss < best_val_loss:
        best_val_loss, no_improve = val_loss, 0
        torch.save({"epoch": epoch, "model": model.state_dict(),
                    "optim": optimizer.state_dict()}, "best_model.pt")
    else:
        no_improve += 1
        if no_improve >= patience:
            break

에폭 루프는 배치 루프를 감싸고, 에폭마다 학습 한 바퀴와 검증 한 바퀴를 돈다. 판단은 검증 손실로 한다. 학습 손실은 에폭을 거듭할수록 거의 항상 내려가므로 아무것도 알려 주지 않고, 모델이 보지 않은 데이터의 손실이 내려가다 올라가기 시작하는 지점이 과적합이 시작되는 자리다. 조기 종료(early stopping)는 그 지점을 자동으로 잡는 장치로, 검증 손실이 patience 에폭 연속으로 최선을 못 넘으면 학습을 멈춘다.

patience를 1로 두지 않는 이유가 있다. 검증 손실은 에폭마다 잡음이 섞여 오르내리고, 한 번 올랐다고 끝내면 다음 에폭에 다시 내려갔을 최선을 놓친다. 다섯 정도면 한두 번의 흔들림은 넘기고 진짜 추세만 잡는다. 대신 멈추는 시점의 모델은 최선이 아니라 최선에서 다섯 에폭 지난 모델이라, 최선을 따로 저장해 두어야 한다. 위 코드가 검증 손실이 갱신될 때만 저장하는 것이 그 때문이고, 학습이 끝나면 마지막 상태가 아니라 best_model.pt를 불러 쓴다.

체크포인트와 재개

저장하는 것이 모델의 state_dict만이 아니라는 점이 이 코드의 요점이다. epoch과 옵티마이저의 state_dict가 함께 들어간다. 체크포인트는 학습을 그 시점부터 다시 이어 갈 수 있게 필요한 상태를 통째로 적어 둔 것이고, 모델 파라미터만으로는 이어 갈 수 없다.

Adam 계열 옵티마이저는 파라미터마다 기울기의 이동 평균 둘을 들고 다닌다. 지난 스텝들의 방향과 크기를 기억해 갱신폭을 조절하는 값이라 모멘텀이라고도 부른다. 이것을 버리고 파라미터만 복원한 뒤 새 옵티마이저로 시작하면, 첫 몇 스텝의 갱신이 학습 초기와 같은 크기로 튀어 애써 다듬은 파라미터를 흔든다. 옵티마이저 상태까지 복원하면 중단 직전과 같은 걸음으로 이어진다. epoch을 적어 두는 것은 재개했을 때 루프 변수와 스케줄러를 그 자리에 맞추기 위해서다.

재개할 때는 저장한 사전을 읽어 model.load_state_dict(ckpt["model"]), optimizer.load_state_dict(ckpt["optim"])으로 각각 채우고 range(ckpt["epoch"] + 1, ...)부터 돈다. 스케줄러가 있으면 그 상태도 같이 저장·복원한다. 긴 학습을 도는 환경에서는 최선 모델과 별개로 매 에폭 끝의 상태를 「마지막 체크포인트」로 덮어써 두는 편이 안전하다. 최선은 되돌아갈 자리이고 마지막은 이어 갈 자리라 쓰임이 다르다.

학습률 스케줄링

학습률 스케줄러

from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau

scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)
scheduler = ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=3)

학습률은 옵티마이저가 기울기 방향으로 한 번에 얼마나 옮길지를 정하는 값이고, 처음엔 크게 걷다가 나중에 잘게 걷는 것이 대개 낫다. 이 변화를 자동으로 주는 것이 학습률 스케줄러다. 갈래가 둘이다. 하나는 시간표대로 줄이는 쪽으로, StepLR은 정해진 에폭마다 일정 비율로 내리고 CosineAnnealingLR은 코사인 곡선을 따라 T_max 에폭에 걸쳐 eta_min까지 부드럽게 내린다. 학습 길이를 미리 정해 두는 이미지 모델 학습에서 많이 쓴다. 다른 하나는 지표를 보고 줄이는 쪽으로, ReduceLROnPlateau는 검증 손실이 patience 에폭 동안 안 내려가면 학습률에 factor를 곱한다. 언제 정체할지 미리 모르는 언어 모델 학습 같은 자리에 맞는다.

앞의 에폭 루프에 둘째 것을 걸어 두었으니 숫자를 맞춰 읽어 본다. 스케줄러의 patience가 3이고 조기 종료의 patience가 5다. 검증 손실이 정체하면 스케줄러가 넷째 정체 에폭에 학습률을 절반으로 내리고, 조기 종료는 다섯째에 멈추므로 낮아진 학습률로 되살아날 기회는 한 에폭뿐이다. 그래도 순서는 맞다 — 스케줄러의 인내가 조기 종료의 인내보다 짧아야 「학습률을 줄여 볼 기회」가 「그만두기」보다 먼저 온다. 두 값을 거꾸로 두면 스케줄러가 한 번도 일하기 전에 학습이 끝나고, 여유를 더 주고 싶으면 조기 종료 쪽 값을 올린다.

step 인자

두 갈래는 부르는 법이 다르다. 시간표대로 가는 스케줄러는 에폭 끝에 scheduler.step()을 인자 없이 부르고, 지표를 보는 스케줄러는 scheduler.step(val_loss)처럼 지표를 넘긴다. 둘을 바꿔 부르면 증상이 다르다. 뒤의 것에 지표를 안 넘기면 필수 인자가 빠졌다는 TypeError가 나서 바로 잡히지만, 앞의 것에 지표를 넘기면 오류 없이 돈다 — 넘긴 값을 에폭 번호로 받아 학습률을 계산하므로 0.43 같은 검증 손실이 「0.43번째 에폭」이 되어 학습률이 엉뚱하게 잡히고, 남는 것은 경고 한 줄뿐이다. 스케줄러를 갈아 끼울 때 step 줄도 함께 봐야 하는 이유다. 부르는 위치는 에폭 루프 안, 그 에폭의 검증이 끝난 자리여야 한다. 배치 루프 안에서 부르면 배치마다 한 에폭치가 줄어 몇 배치 만에 학습률이 바닥에 붙는다.

num_epochs에 맞춰 코사인 곡선을 그리는 스케줄러와 조기 종료는 서로 어긋난다는 점도 적어 둔다. 곡선이 끝나기 전에 멈추면 학습률이 충분히 내려가기 전이라 마지막 미세 조정을 못 한 채 끝나는 셈이다. 그래서 코사인은 정해진 길이를 다 도는 학습에, 정체 감지형은 조기 종료와 함께 쓰는 것이 자연스러운 짝이다.

학습 루프 정리

돌아보면 이 글에서 「학습」 자체는 배치 루프의 다섯 줄이 전부다. 나머지는 그 다섯 줄이 제대로 돌게 하는 장치다. 텐서와 autograd는 backward()가 무엇을 채우는지를 설명하고, nn.Module은 model(x)가 무엇을 부르는지를, DataLoader는 for x, y in loader가 무엇을 내주는지를 설명한다. 검증·조기 종료·체크포인트·스케줄러는 그 루프를 언제 멈추고 무엇을 남기고 걸음을 어떻게 조절할지의 이야기다.

그래서 새 프로젝트를 시작할 때 손이 가는 순서도 이 글의 순서와 같다. Dataset이 샘플 하나를 제대로 내주는지 ds[0]으로 찍어 보고, DataLoader가 배치 모양을 맞게 쌓는지 next(iter(loader))로 보고, 모델에 그 배치를 한 번 통과시켜 출력 모양과 파라미터 수를 확인하고, 배치 하나로 다섯 줄을 몇 번 돌려 손실이 내려가는지 본다. 배치 하나에서 손실이 안 내려가면 데이터를 다 돌려도 안 내려간다. 이 네 단계가 각각 10초짜리 확인이고, 이것을 건너뛰고 전체 학습을 돌리면 한 시간 뒤에 같은 문제를 훨씬 멀리서 만난다.

다음 걸음

이 글의 루프는 전부 손으로 적었다. 배치를 꺼내고, 장치로 옮기고, 다섯 줄을 부르고, 검증하고, 저장하는 것까지 한 줄 한 줄 우리가 썼다. 그래서 어느 줄이 무엇을 하는지 보이고 어느 줄이든 바꿀 수 있다는 것이 PyTorch의 방식이고, 연구 코드가 이쪽으로 모인 이유이기도 하다.

같은 일을 반대편에서 접근하는 프레임워크가 있다. TensorFlow/Keras다. 모델을 층 목록으로 선언하고, 손실과 옵티마이저를 한 번 지정하고, fit 한 번에 에폭 루프와 검증과 콜백까지 맡기는 방식이다. 다음 글에서는 그 고수준 API로 같은 학습을 몇 줄로 줄이는지, 그리고 그 몇 줄 뒤에서 이 글의 다섯 줄이 어떻게 숨어 도는지를 본다. 이 글을 읽었다면 그 fit 한 줄이 무엇을 대신 해 주는지가 보일 것이다.

---

읽어주셔서 감사합니다. 😊