홈 / 강의 / AI 이미지·애니메이션 생성, 원리부터 / 0단계. 최소 기초

0-6. 실습: 작은 디노이저 학습

0단계에서 배운 것을 전부 모아 실제로 돌려 본다. 손글씨 숫자 그림에 노이즈를 섞고, 작은 U-Net이 그 노이즈를 맞히게 가르친다. 코드는 80줄이고 학습은 1분이면 끝난다.

무엇을 만드나

데이터: MNIST. 28×28 크기의 손글씨 숫자 그림 6만 장
모델: 0-4의 U-Net을 아주 작게 만든 것. 파라미터 약 50만 개
학습: 0-1의 식으로 노이즈를 섞고, 섞은 노이즈를 맞히게 한다
확인: 노이즈가 반쯤 섞인 그림에서 원본을 되찾아 본다

실제 이미지 생성 모델의 학습과 뼈대가 같다. 크기만 다르다.

준비

파이썬과 PyTorch를 쓴다. 이미지 생성 모델과 학습 도구가 전부 PyTorch로 만들어져 있어서, 이 분야를 공부하려면 피할 수 없다.

pip install torch torchvision

그래픽카드가 없어도 돌아간다. 다만 몇 분 걸린다.

전체 코드

전체 코드는 0-6-denoiser.py다. 아래에서 부분별로 설명한다.

1. 데이터

tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.5, 0.5)])
data = datasets.MNIST("data", train=True, download=True, transform=tf)
loader = DataLoader(data, batch_size=128, shuffle=True, drop_last=True)

픽셀 값을 −1에서 1 사이로 맞춘다. 0-1에서 본 대로 노이즈 값은 대부분 −3에서 3 사이다. 그림과 노이즈의 크기를 비슷하게 맞춰 놓는 것이다.
batch_size=128: 한 번에 그림 128장을 넣는다. 0-3의 미니배치다.

2. 모델

class Block(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
        self.time = nn.Linear(64, cout)

    def forward(self, x, temb):
        h = F.silu(self.conv1(x))
        h = h + self.time(temb)[:, :, None, None]
        return F.silu(self.conv2(h))

블록 하나다. 3×3 합성곱 두 번이고, 그 사이에 시간 정보를 더한다.

nn.Conv2d(cin, cout, 3, padding=1): 0-4의 합성곱. 채널을 cin에서 cout으로 바꾼다. 그림 크기는 그대로다.
F.silu: 0-3의 활성화 함수.
self.time(temb): 노이즈 비율 정보를 채널 수에 맞게 바꿔서 더한다.

class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.tmlp = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
        self.down1 = Block(1, 32)         # 28x28
        self.down2 = Block(32, 64)        # 14x14
        self.mid = Block(64, 128)         # 7x7
        self.up2 = Block(128 + 64, 64)    # 14x14
        self.up1 = Block(64 + 32, 32)     # 28x28
        self.out = nn.Conv2d(32, 1, 1)

    def forward(self, x, t):
        temb = self.tmlp(t[:, None])
        d1 = self.down1(x, temb)
        d2 = self.down2(F.avg_pool2d(d1, 2), temb)
        m = self.mid(F.avg_pool2d(d2, 2), temb)
        u2 = self.up2(torch.cat([F.interpolate(m, scale_factor=2), d2], dim=1), temb)
        u1 = self.up1(torch.cat([F.interpolate(u2, scale_factor=2), d1], dim=1), temb)
        return self.out(u1)

0-4의 U자 그림을 그대로 옮긴 것이다.

tmlp: 숫자 하나(t)를 숫자 64개로 바꾼다. 시간 임베딩이다.
avg_pool2d(…, 2): 절반으로 줄인다. 28 → 14 → 7.
interpolate(…, scale_factor=2): 두 배로 키운다. 7 → 14 → 28.
torch.cat([…, d2], dim=1): 스킵 연결. 내려갈 때의 결과 d2를 올라올 때 옆에 붙인다. 그래서 up2의 입력 채널이 128 + 64다.
self.out: 마지막에 채널을 1개로 줄인다. 입력과 같은 크기의 노이즈 지도가 나온다.

3. 학습

for epoch in range(3):
    for x0, _ in loader:
        x0 = x0.to(device)
        t = torch.rand(x0.shape[0], device=device)
        eps = torch.randn_like(x0)
        tt = t[:, None, None, None]
        xt = (1 - tt) * x0 + tt * eps
        loss = F.mse_loss(model(xt, t), eps)
        opt.zero_grad()
        loss.backward()
        opt.step()

핵심은 가운데 다섯 줄이다.

코드	뜻
`t = torch.rand(…)`	그림마다 노이즈 비율을 0~1에서 무작위로 고른다
`eps = torch.randn_like(x0)`	픽셀마다 따로 노이즈를 뽑는다. 0-1의 ε
`xt = (1 - tt) * x0 + tt * eps`	0-1의 섞는 식 그대로
`model(xt, t)`	모델이 섞인 그림과 t를 보고 노이즈를 예측한다
`F.mse_loss(…, eps)`	예측과 실제 노이즈의 제곱 오차. 0-2, 0-3의 손실

아래 세 줄은 0-3의 학습 반복이다. loss.backward()가 역전파, opt.step()이 파라미터를 조금 움직이는 일이다.

정답을 따로 만들 필요가 없다는 점을 눈여겨보자. 노이즈를 우리가 직접 섞었으니 정답(ε)을 이미 알고 있다. 그림만 있으면 학습 문제를 끝없이 만들어 낼 수 있다.

4. 확인

xt = (1 - tt) * x0 + tt * eps
with torch.no_grad():
    eps_hat = model(xt, t)
x0_hat = (xt - tt * eps_hat) / (1 - tt)

섞는 식을 x₀에 대해 풀면 x₀ = (x_t − t·ε) / (1 − t)다. 진짜 ε 자리에 모델이 예측한 값을 넣으면 원본을 추정할 수 있다.

돌려 본 결과

파라미터 수: 496065
에폭 1: 평균 손실 0.1949
에폭 2: 평균 손실 0.0850
에폭 3: 평균 손실 0.0724
복원 오차: 0.0464

손실이 내려간다. 아무것도 모르는 모델이 0만 내놓으면 손실이 1 근처다(노이즈의 분산이 1이므로). 0.07이면 노이즈의 대부분을 맞히고 있다는 뜻이다.

윗줄: 원본
가운뎃줄: 노이즈를 반(t = 0.5) 섞은 것
아랫줄: 모델이 되찾은 원본

여기서 한 것과 안 한 것

한 것: 노이즈가 섞인 그림에서 노이즈를 맞히는 모델을 만들었다. 이것이 디노이저다. 이미지 생성 모델의 심장이다.

안 한 것: 순수 노이즈에서 출발해 새 그림을 만들어 내지는 않았다. 디노이저를 여러 번 불러서 노이즈를 조금씩 걷어 내야 하는데, 그 방법이 1단계의 주제다.

해 볼 것

확인 코드에서 t를 0.2, 0.8로 바꿔 보자. 얼마나 섞였을 때까지 되찾을 수 있나?
스킵 연결을 빼 보자. torch.cat을 지우고 up2, up1의 입력 채널을 128, 64로 고치면 된다. 결과가 어떻게 달라지나?
에폭을 10으로 늘려 보자. 손실이 어디까지 내려가나?

예상 결과

t = 0.2면 거의 완벽하게 되찾는다. t = 0.8이면 숫자의 대략적인 모양만 나오고, 다른 숫자로 바뀌기도 한다. 정보가 그만큼 사라졌기 때문이다.
손실이 더 높은 곳에서 멈추고 결과가 뭉개진다. 0-4에서 말한 스킵 연결의 역할이다.
계속 조금씩 내려가지만 0이 되지는 않는다. 노이즈가 많이 섞인 그림에서는 정답을 정확히 알 수 없기 때문이다. 모델은 가능한 답들의 평균을 낸다.