모두를 위한 AI 5편 — 47만 파라미터 U-Net으로 흑백 사진에 색 입히기, 그리고 왜 색이 바랬는가

• 들어가며 — 출력이 이미지일 때
• U-Net — skip connection이 나르는 것
• 학습
• 결과 — 형태는 맞고 색은 바랬다
• 왜 색이 바랬나 — L1 손실의 성질
• 실제 컬러화 모델은 어떻게 하는가
• 정리
• 🧠 이해도 체크 퀴즈
• 참고 자료

들어가며 — 출력이 이미지일 때

지금까지는 출력이 텍스트였습니다. 1편은 문장, 3편은 단어, 4편은 캡션이었죠. 이번 편은 출력 자체가 이미지입니다.

문제는 컬러화입니다. 흑백 사진을 넣으면 색을 입혀 내놓습니다.

입력: 32×32 흑백 (채널 1개)
출력: 32×32 컬러 (채널 3개)

데이터는 따로 구할 필요가 없습니다. CIFAR-10의 컬러 이미지를 흑백으로 바꾸면 입력과 정답 쌍이 저절로 만들어집니다.

rgb = torch.stack([...])                                    # 원본 컬러 = 정답
gray = (rgb * torch.tensor([0.299, 0.587, 0.114])
        [None, :, None, None]).sum(1, keepdim=True)         # 흑백 변환 = 입력

가중치 0.299, 0.587, 0.114는 사람 눈이 초록에 가장 민감하고 파랑에 둔감하다는 점을 반영한 표준 휘도 공식입니다. 이렇게 만든 데이터를 자기지도 학습(self-supervised)이라고 부릅니다. 사람이 라벨을 붙이지 않았는데 정답이 존재하니까요.

U-Net — skip connection이 나르는 것

컬러화에는 까다로운 요구가 있습니다. "무엇인지" 알아야 색을 정하는데, "어디인지"도 정확해야 합니다. 개구리인 줄 알아야 초록을 칠하지만, 개구리의 윤곽선 바깥으로 초록이 새면 안 됩니다.

일반적인 인코더-디코더는 이 둘을 동시에 만족시키기 어렵습니다. 인코더가 이미지를 압축하면서 의미는 얻지만 위치 정보를 잃기 때문입니다. U-Net은 이를 skip connection으로 해결합니다.

class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.e1 = block(1, 32); self.e2 = block(32, 64); self.e3 = block(64, 128)
        self.d2 = block(128 + 64, 64); self.d1 = block(64 + 32, 32)
        self.out = nn.Conv2d(32, 3, 1)

    def forward(self, x):
        s1 = self.e1(x)                       # 32x32  — 가장 세밀
        s2 = self.e2(F.max_pool2d(s1, 2))     # 16x16
        h  = self.e3(F.max_pool2d(s2, 2))     # 8x8    — 가장 추상적
        h = self.d2(torch.cat([F.interpolate(h, scale_factor=2), s2], 1))
        h = self.d1(torch.cat([F.interpolate(h, scale_factor=2), s1], 1))
        return torch.sigmoid(self.out(h))

핵심은 torch.cat([업샘플된_h, s2], 1) 입니다. 디코더가 해상도를 되돌릴 때, 같은 해상도의 인코더 출력을 그대로 가져와 이어 붙입니다.

• 아래에서 올라온 h — 저해상도지만 "이건 개구리다" 같은 의미를 담고 있음
• 옆에서 온 s2 — 압축을 거치지 않아 가장자리와 질감이 온전함

디코더는 이 둘을 함께 봅니다. 의미는 깊은 경로에서, 위치는 skip 경로에서 옵니다. 채널 수가 128 + 64 인 이유가 이것입니다.

sigmoid로 끝나는 것도 의도적입니다. 출력이 0~1 범위의 RGB 값이어야 하니 활성화로 범위를 강제합니다.

학습

opt = torch.optim.AdamW(model.parameters(), lr=2e-3)

while not run.over_budget():          # 12분 예산
    ix = torch.randint(0, len(rgb), (128,))
    x, y = gray[ix].to(dev), rgb[ix].to(dev)
    loss = F.l1_loss(model(x), y)
    opt.zero_grad(); loss.backward(); opt.step()

손실은 L1(평균 절대 오차)입니다. 예측한 RGB 값과 정답 RGB 값의 차이를 그대로 재는 회귀 손실입니다. 이 선택이 결과의 성격을 결정하는데, 뒤에서 다시 다루겠습니다.

[   14.1s] pairs=30000
[   14.2s] params {"params": 472323, "res": 32}
[  722.0s] SUMMARY {"steps": 62136, "final_l1": 0.0231}

62,136 스텝, 최종 L1 0.023. 픽셀 값이 0~1 범위이니 평균 오차가 2.3% 라는 뜻입니다. 숫자만 보면 훌륭합니다.

결과 — 형태는 맞고 색은 바랬다

각 행은 세 장씩 두 묶음이고, 순서는 입력 흑백 / 모델 예측 / 정답 컬러 입니다.

정직하게 읽어 봅시다.

잘 된 것 — 형태가 정확히 보존됩니다. 고양이의 털, 배의 갑판 구조, 개구리의 다리, 자동차의 창틀이 흐려지지 않았습니다. skip connection이 제 역할을 했다는 증거입니다. 색의 방향도 대체로 맞습니다. 개구리 주변은 초록빛, 배 주변은 푸른빛, 고양이는 갈색 계열로 갔습니다.

안 된 것 — 색이 전반적으로 바랬습니다. 정답에서 선명한 빨강인 자동차가 예측에서는 어두운 회갈색입니다. 정답의 붉은 선체와 파란 바다가 예측에서는 탁한 회청색입니다. 세피아 톤에 가까운 결과가 많습니다.

왜 색이 바랬나 — L1 손실의 성질

이건 모델이 작아서가 아닙니다. 손실 함수를 L1으로 고른 결과입니다.

흑백 자동차 사진 한 장을 생각해 봅시다. 이 차는 빨강일 수도, 파랑일 수도, 흰색일 수도 있습니다. 흑백 정보만으로는 결정할 수 없습니다. 즉 정답이 하나가 아니라 여러 개인 문제입니다.

이때 L1 손실은 모델에게 이렇게 말합니다. "정답과의 절대 오차를 최소화하라." 여러 정답이 가능한 상황에서 이 지시를 따르는 가장 안전한 방법은 무엇일까요.

모든 가능성의 중간값을 찍는 것입니다.

빨강을 찍었다가 정답이 파랑이면 큰 벌점을 받습니다. 하지만 회색을 찍으면 정답이 무엇이든 벌점이 중간 정도로 억제됩니다. 학습이 진행될수록 모델은 이 안전한 전략으로 수렴하고, 그 결과가 채도 낮은 색입니다.

L1 0.023이라는 좋은 숫자와 바랜 색이 동시에 성립하는 이유가 이것입니다. L1은 정확히 자기가 시킨 일을 잘한 것이고, 우리가 원한 것이 그게 아니었을 뿐입니다.

실제 컬러화 모델은 어떻게 하는가

이 문제는 잘 알려져 있고, 해법도 여러 갈래로 나와 있습니다.

색을 분류 문제로 바꾸기 — Zhang 등의 2016년 연구는 색 공간을 313개 구간으로 나누고 "이 픽셀이 어느 구간인가"를 맞히는 분류로 풀었습니다. 분류에서는 여러 후보에 확률을 나눠 줄 수 있으므로 평균으로 뭉개지지 않습니다. 샘플링할 때 선명한 색을 고르면 됩니다.

적대적 손실 추가 — GAN의 판별자는 "이 이미지가 진짜 같은가"를 봅니다. 바랜 색은 진짜 사진처럼 보이지 않으므로 벌점을 받습니다. L1이 형태를, 판별자가 선명도를 담당하는 조합입니다.

지각 손실 사용 — 픽셀 값을 직접 비교하는 대신, 사전학습된 신경망의 특징 공간에서 비교합니다. 사람이 느끼는 유사성에 더 가깝습니다.

셋 다 공통점이 있습니다. 평균으로 도망가는 것이 이득이 되지 않도록 손실을 설계한다는 점입니다.

정리

47만 파라미터로 형태를 온전히 보존하며 색을 입혔습니다. skip connection 덕분에 세밀한 구조가 살았고, 자기지도 방식이라 라벨링 비용이 0이었습니다.

그리고 이 편의 진짜 교훈은 색이 바랜 쪽에 있습니다. 손실 함수는 목표를 정의하는 것이지 목표를 달성하는 도구가 아닙니다. L1을 고른 순간 "평균에 가까운 답이 안전하다"는 규칙이 함께 정해졌고, 모델은 그 규칙을 충실히 따랐습니다. 결과가 마음에 들지 않는다면 모델이 아니라 무엇을 최소화하라고 시켰는지를 먼저 돌아봐야 합니다.

다음 편에서는 이 문제를 정면으로 다루는 방식 — 디퓨전을 봅니다. 여러 정답이 가능한 상황에서 평균으로 도망가지 않고 하나를 골라내는 구조입니다.

🧠 이해도 체크 퀴즈

참고 자료

• CIFAR-10 데이터셋
• U-Net — 생물의학 이미지 분할을 위한 합성곱 신경망
• Colorful Image Colorization — 색을 분류로 푸는 접근
• Image-to-Image Translation with Conditional Adversarial Networks (pix2pix)