SwinIR 도트 업스케일 학습 — 한 파일로 따라가기

학습 코드 4개(swinir_finetune/data.py, swinir_finetune/learning.py, swinir_point/outline.py, swinir_point/data.py)를 그대로 옮겨 붙이고, 각 단계의 중간 결과를 그림과 print로 확인합니다. 마지막에는 실제 10,000스텝 대신 40스텝만 짧게 학습해 손실이 내려가고 결과가 바뀌는 것을 눈으로 봅니다.

실행 환경: 커널 Python (ComfyUI venv · maple) · GPU 1장 · 프로젝트 루트 ~/prj/maple/6_upscale

0. 준비 — 경로·라이브러리·한글 글꼴, 데이터 개수 확인

import sys, json, time, hashlib
from pathlib import Path
from functools import lru_cache
from collections import defaultdict
import numpy as np, torch
from PIL import Image
import matplotlib; import matplotlib.pyplot as plt
from matplotlib import font_manager

ROOT = Path.home()/'prj/maple/6_upscale'; sys.path.insert(0, str(ROOT))
font_manager.fontManager.addfont('/usr/share/fonts/truetype/nanum/NanumSquareRoundB.ttf'); plt.rcParams['font.family'] = 'NanumSquareRound'; plt.rcParams['axes.unicode_minus'] = False
DATA = ROOT/'data/swinir_finetune_v1'                                    # 검증된 학습 쌍 캐시(HR 192·LR 96·마스크·투명 원본)
WEIGHTS = Path('/home/son/data_ssd/public/models/maple_upscale/pretrained/weights/001_classicalSR_DF2K_s64w8_SwinIR-M_x2.pth')
summary = json.loads((DATA/'summary.json').read_text())
print('torch', torch.__version__, '| GPU', torch.cuda.get_device_name(0))
print('학습 쌍 장수:', summary['split_counts'])

def show(images, titles, scale=1, size=2.6, cmap=None):
    # 이미지 여러 장을 한 줄에 놓는다. 도트가 보이게 보간 없이(nearest) 그린다.
    fig, axes = plt.subplots(1, len(images), figsize=(size*len(images), size+0.5))
    for ax, im, t in zip(np.atleast_1d(axes), images, titles):
        ax.imshow(im, interpolation='nearest', cmap=cmap); ax.set_title(t, fontsize=10); ax.axis('off')
    plt.tight_layout(); plt.show()

def to_img(t):
    # 3×H×W [0,1] 텐서 → H×W×3 uint8
    return t.detach().clamp(0, 1).permute(1, 2, 0).mul(255).round().byte().cpu().numpy()

/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:65: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]

torch 2.10.0+cu128 | GPU NVIDIA RTX PRO 6000 Blackwell Workstation Edition
학습 쌍 장수: {'test': 274, 'train': 2246, 'val': 262}

1. 데이터 읽기 — swinir_finetune/data.py 그대로

학습용은 train·val만 열 수 있고, 파일 해시가 목록과 다르면 멈춥니다(test·옛 도트가 섞이는 사고 방지).

def digest(path):
    # 파일 바이트를 기준으로 캐시와 원본의 변경을 확인합니다.
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def manifest_rows(scale, split, directory=DATA):
    # 학습용 경로는 train과 val만 열고 최종 시험 및 옛 도트는 차단합니다.
    if scale not in (2, 3, 4) or split not in ("train", "val"):
        raise ValueError("학습 데이터는 scale=2/3/4, split=train/val만 허용합니다.")
    directory = Path(directory)
    report = json.loads((directory / "summary.json").read_text())
    manifest = directory / "manifest.jsonl"
    assert report["status"] == "PASS" and digest(manifest) == report["manifest_sha256"]
    assert report["split_family_overlaps"] == report["crop_split_overlaps"] == 0
    return [row for line in manifest.read_text().splitlines()
            if (row := json.loads(line))["split"] == split]


class Patches(torch.utils.data.Dataset):
    # 동일한 HR에서 만든 세 배율의 LR을 모델별로 선택합니다.
    def __init__(self, scale, split="train", directory=DATA):
        self.scale, self.split, self.directory = scale, split, Path(directory)
        self.rows = manifest_rows(scale, split, directory)

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

    @lru_cache(maxsize=256)
    def __getitem__(self, index):
        # PNG 해시와 크기를 확인한 뒤 RGB 및 전경 마스크를 [0,1]로 반환합니다.
        row = self.rows[index]
        result = {"index": index}
        for key, stored in (("lr", f"lr_x{self.scale}"), ("hr", "hr"), ("mask", "mask")):
            path = self.directory / row["files"][stored]
            assert digest(path) == row["sha256"][stored], f"캐시가 변경됐습니다: {path}"
            with Image.open(path) as image:
                raw = np.array(image.convert("L" if key == "mask" else "RGB"), copy=True)
            tensor = torch.from_numpy(raw).float().div_(255)
            result[key] = tensor.unsqueeze(0) if key == "mask" else tensor.permute(2, 0, 1)
        assert result["hr"].shape == (3, 192, 192) and result["mask"].shape == (1, 192, 192)
        assert result["lr"].shape == (3, 192 // self.scale, 192 // self.scale)
        return result

train, val = Patches(2, 'train'), Patches(2, 'val')
print('train', len(train), '장 / val', len(val), '장 · 캐릭터 수 train', len({r['family'] for r in train.rows}), 'val', len({r['family'] for r in val.rows}))
try:
    Patches(2, 'test')
except ValueError as e:
    print('test를 열면 →', e)
item = train[0]
print('한 항목의 텐서 모양:', {k: tuple(v.shape) for k, v in item.items() if k != 'index'}, '| 값 범위 0~1, 이름:', train.rows[0]['name'])

train 2246 장 / val 262 장 · 캐릭터 수 train 400 val 50
test를 열면 → 학습 데이터는 scale=2/3/4, split=train/val만 허용합니다.
한 항목의 텐서 모양: {'lr': (3, 96, 96), 'hr': (3, 192, 192), 'mask': (1, 192, 192)} | 값 범위 0~1, 이름: 쿠보

# 학습 쌍 4개: HR 정답(192) · 캐시의 Bicubic LR(96, 2배로 키워 표시) · 전경 마스크
for i in (0, 300, 900, 1500):
    it = train[i]; r = train.rows[i]
    show([to_img(it['hr']), np.kron(to_img(it['lr']), np.ones((2, 2, 1), np.uint8)), it['mask'][0].numpy()],
         [f"HR 192×192 · {r['name']}", 'Bicubic LR 96×96 (첫 실험 입력)', '전경 마스크 alpha>0'], cmap='gray')

2. 경계선 보존 점 샘플링 — swinir_point/outline.py 그대로

Bicubic은 색을 섞어 흐리게 만듭니다. 대신 2×2 블록마다 HR 픽셀 하나를 그대로 뽑되, 검은 윤곽 픽셀이 있으면 그것을 우선 살립니다.

DARK = 70          # 평균 밝기(R+G+B)/3 가 이 값 이하이면 '어두운 픽셀' (실측: 검정·17단계 회색·어두운 갈색)
CONTRAST = 80      # 이웃과 밝기 차이가 이 값 이상이면 '경계'


def outline_mask(rgba):
    # 어두운 픽셀 중 상하좌우 이웃에 훨씬 밝은 픽셀이나 투명 픽셀이 있는 것만 경계선으로 본다.
    rgb, alpha = rgba[..., :3].astype(np.int16), rgba[..., 3] > 0
    lum = rgb.sum(-1)//3
    dark = (lum <= DARK) & alpha
    pad = np.pad(lum, 1, constant_values=255); pada = np.pad(alpha, 1, constant_values=False)
    edge = np.zeros_like(dark)
    for dy, dx in ((0, 1), (0, -1), (1, 0), (-1, 0)):
        nl = pad[1+dy:1+dy+lum.shape[0], 1+dx:1+dx+lum.shape[1]]
        na = pada[1+dy:1+dy+lum.shape[0], 1+dx:1+dx+lum.shape[1]]
        edge |= (nl - lum >= CONTRAST) | ~na
    return dark & edge, dark, alpha, lum


def outline_sample(rgba, scale=2):
    # 블록마다 우선순위: 경계 픽셀 > 어두운 픽셀 > 불투명 픽셀 > 기본 위치 HR[2i+1,2j+1].
    # 같은 급이 여럿이면 기본 위치(오른쪽 아래)에 가장 가까운 것을 골라 출력 윤곽의 반 픽셀 흔들림을 줄인다.
    h, w = rgba.shape[0]//scale*scale, rgba.shape[1]//scale*scale
    rgba = rgba[:h, :w]
    edge, dark, alpha, lum = outline_mask(rgba)
    def blocks(x): return x.reshape(h//scale, scale, w//scale, scale, *x.shape[2:]).swapaxes(1, 2).reshape(h//scale, w//scale, scale*scale, *x.shape[2:])
    b_rgba, b_edge, b_dark, b_alpha, b_lum = (blocks(v) for v in (rgba, edge, dark, alpha, lum))
    default = scale*(scale//2) + scale//2                     # 2×2에서 인덱스 3 = (1,1)
    ys, xs = np.divmod(np.arange(scale*scale), scale)
    proximity = -(np.abs(ys - scale//2) + np.abs(xs - scale//2))            # 기본 위치와의 거리
    score = np.where(b_edge, 3000, 0) + np.where(b_dark, 2000, 0) + np.where(b_alpha, 1000, 0) + proximity*10
    choice = score.argmax(-1)
    lr = np.take_along_axis(b_rgba, choice[..., None, None], 2)[:, :, 0]
    stats = dict(hr_outline_pixels=int(edge.sum()), lr_outline_pixels=int(np.take_along_axis(b_edge, choice[..., None], 2).sum()),
                 blocks_with_outline=int(b_edge.any(-1).sum()), blocks_changed=int((choice != default).sum()), blocks=int(choice.size))
    return lr, stats


def point_lr_rgb(rgb, scale=2):
    # 비교용 기본 점 샘플링(HR[2i+1, 2j+1]).
    r = scale//2
    return rgb[r::scale, r::scale]


def white_rgb(rgba):
    # RGBA(uint8) → 흰 배경 합성 RGB(uint8). 학습 HR과 같은 규칙.
    a = rgba[..., 3:4].astype(np.float32)/255
    return (rgba[..., :3].astype(np.float32)*a + 255*(1-a)).round().astype(np.uint8)

i = 300; r = train.rows[i]
rgba = np.asarray(Image.open(DATA/r['files']['rgba']).convert('RGBA'))          # 투명 원본(192×192 RGBA)
hr = white_rgb(rgba); edge, dark, alpha, lum = outline_mask(rgba)
lr_point = point_lr_rgb(hr); lr_out_rgba, st = outline_sample(rgba); lr_out = white_rgb(lr_out_rgba)
lr_bic = to_img(train[i]['lr'])
print(r['name'], '| HR 경계 픽셀', st['hr_outline_pixels'], '→ 점 샘플링에 남는 것', int(edge[1::2, 1::2].sum()), '/ 경계선 보존에 남는 것', st['lr_outline_pixels'],
      f"| 블록 {st['blocks_changed']}/{st['blocks']}에서 기본 위치 대신 다른 픽셀 선택")
up = lambda a: np.kron(a, np.ones((2, 2, 1), np.uint8))
show([hr, edge, up(lr_bic), up(lr_point), up(lr_out)], ['HR 192', '경계 픽셀 마스크', 'Bicubic LR (흐림)', '점 샘플링 LR HR[2i+1,2j+1]', '경계선 보존 LR'], cmap='gray')

블랙윙 보안담당자 | HR 경계 픽셀 1382 → 점 샘플링에 남는 것 338 / 경계선 보존에 남는 것 777 | 블록 480/9216에서 기본 위치 대신 다른 픽셀 선택

# 같은 자리를 크게: 24×24 HR 영역과, 거기서 나온 12×12 LR 세 종류(각각 4배로 키워 표시)
ys_, xs_ = np.where(edge); cy, cx = int(np.median(ys_)), int(np.median(xs_)); y0, x0 = max(cy-12, 0)//2*2, max(cx-12, 0)//2*2
z = lambda a, k: np.kron(a, np.ones((k, k, 1), np.uint8))
show([z(hr[y0:y0+24, x0:x0+24], 4), z(lr_bic[y0//2:y0//2+12, x0//2:x0//2+12], 8), z(lr_point[y0//2:y0//2+12, x0//2:x0//2+12], 8), z(lr_out[y0//2:y0//2+12, x0//2:x0//2+12], 8)],
     ['HR 24×24 (윤곽 근처)', 'Bicubic LR 12×12', '점 샘플링 LR 12×12', '경계선 보존 LR 12×12'], size=3)

3. 점 샘플링·경계선 보존 데이터셋 — swinir_point/data.py + variants.OutlinePatches

캐시의 Bicubic LR 대신, HR(또는 투명 원본)에서 즉석에서 LR을 만듭니다. HR·마스크·분리 규칙은 그대로입니다.

def point_sample(hr, scale=2):
    # HR 텐서(3×H×W)에서 scale칸마다 픽셀 하나를 그대로 뽑는다. 색이 섞이지 않는다.
    r = scale//2
    return hr[:, r::scale, r::scale].contiguous()


class PointPatches(Patches):
    # lr만 점 샘플링으로 바꾸고 hr·mask·분리 규칙은 그대로 둔다.
    def __getitem__(self, index):
        item = dict(super().__getitem__(index))
        item['lr'] = point_sample(item['hr'], self.scale)
        return item


class OutlinePatches(Patches):
    # 캐시 rgba.png에서 경계선 보존 LR을 만들고 흰 배경에 합성한다. hr·mask는 그대로. (채택한 방식)
    def __getitem__(self, index):
        item = dict(super().__getitem__(index))
        rgba = np.asarray(Image.open(self.directory/self.rows[index]['files']['rgba']).convert('RGBA'))
        lr_rgba, _ = outline_sample(rgba, self.scale)
        item['lr'] = torch.from_numpy(white_rgb(lr_rgba)).float().div_(255).permute(2, 0, 1).contiguous()
        return item

train_o, val_o = OutlinePatches(2, 'train'), OutlinePatches(2, 'val')
a, b, c = train[900]['lr'], PointPatches(2, 'train')[900]['lr'], train_o[900]['lr']
print('세 데이터셋의 LR 모양은 모두', tuple(c.shape), '| Bicubic과 경계선 보존 LR의 평균 절대 차이:', f"{(a-c).abs().mean()*255:.1f} / 255")
show([up(to_img(a)), up(to_img(b)), up(to_img(c))], ['Patches (Bicubic)', 'PointPatches', 'OutlinePatches (채택)'])

세 데이터셋의 LR 모양은 모두 (3, 96, 96) | Bicubic과 경계선 보존 LR의 평균 절대 차이: 12.8 / 255

4. 배치 뽑기 · 증강 · 손실 — swinir_finetune/learning.py 그대로

(1) 프레임이 많은 캐릭터가 학습을 지배하지 않게 캐릭터를 먼저 뽑고 그 안에서 한 장, (2) 회전·반전만 LR·HR·마스크에 똑같이, (3) 전경 L1 + 0.1×배경 L1.

class FamilyBatches:
    # 프레임이 많은 한 캐릭터가 학습을 지배하지 않게 캐릭터부터 뽑는다.
    def __init__(self, rows, config, start=0):
        # 단계 번호로 난수를 정하므로 중단 후에도 같은 순서를 재현한다.
        groups = defaultdict(list)
        for i, row in enumerate(rows):
            groups[row['family']].append(i)
        self.groups = [groups[k] for k in sorted(groups)]
        self.config, self.start = config, start

    def __iter__(self):
        # 끝이 정해진 배치만 만들고 선택한 캐릭터 안에서는 프레임을 균등 추출한다.
        for step in range(self.start, self.config['total_steps']):
            rng = np.random.default_rng(np.random.SeedSequence([self.config['data_seed'], step]))
            yield [int(rng.choice(self.groups[g])) for g in
                   rng.integers(len(self.groups), size=self.config['batch_size'])]


def loader(dataset, config, start=0):
    # CPU에서 PNG를 읽는 동안 GPU가 다음 배치를 처리할 수 있게 한다.
    return torch.utils.data.DataLoader(dataset, batch_sampler=FamilyBatches(dataset.rows, config, start),
                      num_workers=config['workers'], pin_memory=True,
                      generator=torch.Generator().manual_seed(config['data_seed']),
                      persistent_workers=config['workers'] > 0)


def augment(batch, step, seed):
    # LR·HR·전경 마스크를 같은 방향으로 돌리며 RGB 색 자체는 바꾸지 않는다.
    rng = np.random.default_rng(np.random.SeedSequence([seed, step, 97]))
    turns, flip = int(rng.integers(4)), bool(rng.integers(2))
    values = {}
    for key in ('lr', 'hr', 'mask'):
        value = batch[key].to('cuda', non_blocking=True)
        value = torch.rot90(value, turns, (-2, -1))
        values[key] = value.flip(-1) if flip else value
    return values


def train_step(model, optimizer, batch, config, step):
    # 넓은 흰 여백보다 원본 도트의 픽셀 오차를 더 중요하게 줄인다.
    model.train()
    values = augment(batch, step, config['seed'])
    optimizer.zero_grad(set_to_none=True)
    prediction = model(values['lr'])
    error = (prediction-values['hr']).abs()
    mask = (values['mask'] > 0).float()
    foreground = (error*mask).sum()/(mask.sum()*3).clamp_min(1)
    background = (error*(1-mask)).sum()/((1-mask).sum()*3).clamp_min(1)
    loss = foreground + 0.1*background
    assert torch.isfinite(loss), 'non-finite training loss'
    loss.backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config['grad_clip'], error_if_nonfinite=True)
    optimizer.step()
    return dict(loss=float(loss.detach()), foreground_l1=float(foreground.detach()),
                background_l1=float(background.detach()), gradient_norm=float(norm))

config = dict(total_steps=40, batch_size=8, data_seed=73191, seed=20260912, workers=0, grad_clip=1.0,
              learning_rate=2e-5, min_learning_rate=2e-6, betas=(0.9, 0.99))   # 실제 실험은 total_steps=10000, workers=2, 나머지 동일
sampler = FamilyBatches(train_o.rows, config)
first = next(iter(sampler)); names = [train_o.rows[i]['name'] for i in first]
print('첫 배치 8장의 캐릭터:', names)
# 캐릭터 균등 추출 확인: 프레임이 8장인 캐릭터와 1장인 캐릭터가 뽑히는 횟수가 비슷해야 한다
cnt = defaultdict(int); frames = {f: len(g) for f, g in zip(sorted({r['family'] for r in train_o.rows}), sampler.groups)}
for step, idx in zip(range(2000), FamilyBatches(train_o.rows, dict(config, total_steps=2000))):
    for i in idx: cnt[train_o.rows[i]['family']] += 1
by_frames = defaultdict(list)
for f, n in frames.items(): by_frames[n].append(cnt[f])
print('프레임 수별 캐릭터당 평균 선택 횟수 (2,000배치):', {n: round(np.mean(v), 1) for n, v in sorted(by_frames.items())})

첫 배치 8장의 캐릭터: ['램나나', '더 수상한 도적', '푸른 켄타우로스', '기사의 창병2', '심사위원단', '머쉬킹', '상황판', '도도']

프레임 수별 캐릭터당 평균 선택 횟수 (2,000배치): {1: np.float64(40.2), 2: np.float64(40.8), 3: np.float64(39.4), 4: np.float64(38.6), 5: np.float64(42.0), 6: np.float64(37.2), 7: np.float64(39.0), 8: np.float64(40.3)}

# 증강 8가지: 같은 학습 쌍의 LR·HR·마스크가 같은 방향으로 돈다
it = train_o[first[0]]
fig, axes = plt.subplots(3, 8, figsize=(16, 6.4))
for k in range(8):
    turns, flip = k % 4, k >= 4
    for row, key in enumerate(('hr', 'lr', 'mask')):
        v = torch.rot90(it[key], turns, (-2, -1)); v = v.flip(-1) if flip else v
        axes[row, k].imshow(to_img(v) if key != 'mask' else v[0].numpy(), interpolation='nearest', cmap='gray'); axes[row, k].axis('off')
    axes[0, k].set_title(f"회전 {turns*90}°{' + 반전' if flip else ''}", fontsize=9)
axes[0, 0].set_ylabel('HR'); plt.suptitle('증강: 회전 4 × 좌우 반전 2 (색은 그대로) · 위 HR / 가운데 LR / 아래 마스크'); plt.tight_layout(); plt.show()

5. 모델 불러오기와 손실 한 번 계산

공개 SwinIR-M x2 가중치를 그대로 불러와 첫 배치에 넣고, 전경 L1·배경 L1·최종 손실을 찍어 봅니다.

from experiments.public_models.multiscale.swinir import load_model
torch.backends.cuda.matmul.allow_tf32 = torch.backends.cudnn.allow_tf32 = False; torch.set_float32_matmul_precision('highest')
torch.manual_seed(config['seed'])
model, n_tensors = load_model(2, WEIGHTS); model.cuda()
print('불러온 텐서 수', n_tensors, '| 파라미터', f"{sum(p.numel() for p in model.parameters())/1e6:.2f}M", '| FP32, TF32 끔')
batch = next(iter(loader(train_o, config)))
with torch.no_grad():
    v = augment(batch, 0, config['seed']); model.eval(); pred = model(v['lr'])
    err = (pred-v['hr']).abs(); m = (v['mask'] > 0).float()
    fg = (err*m).sum()/(m.sum()*3); bg = (err*(1-m)).sum()/((1-m).sum()*3)
print(f"학습 전 · 전경 L1 {fg:.4f} · 배경 L1 {bg:.4f} · 손실 = 전경 + 0.1×배경 = {fg+0.1*bg:.4f}  (값은 0~1 척도, ×255가 8비트 오차)")
print('배치 모양: LR', tuple(v['lr'].shape), '→ 출력', tuple(pred.shape), '· HR', tuple(v['hr'].shape))
show([up(to_img(v['lr'][0])), to_img(pred[0]), to_img(v['hr'][0]), (err[0].mean(0)*4).clamp(0, 1).cpu().numpy()],
     ['입력 LR (경계선 보존)', '학습 전 출력', 'HR 정답', '|출력−정답| (밝을수록 큼)'], cmap='magma')

/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/requests/__init__.py:113: RequestsDependencyWarning: urllib3 (2.6.3) or chardet (7.4.0.post2)/charset_normalizer (3.4.4) doesn't match a supported version!
  warnings.warn(

/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/timm/models/layers/__init__.py:49: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning)
/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/torch/functional.py:505: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /pytorch/aten/src/ATen/native/TensorShape.cpp:4381.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]

불러온 텐서 수 548 | 파라미터 11.75M | FP32, TF32 끔

학습 전 · 전경 L1 0.1393 · 배경 L1 0.0211 · 손실 = 전경 + 0.1×배경 = 0.1414  (값은 0~1 척도, ×255가 8비트 오차)
배치 모양: LR (8, 3, 96, 96) → 출력 (8, 3, 192, 192) · HR (8, 3, 192, 192)

/tmp/ipykernel_1051820/1034973226.py:23: UserWarning: Glyph 8722 (\N{MINUS SIGN}) missing from font(s) NanumSquareRound.
  plt.tight_layout(); plt.show()
/home/son/prj/ComfyUI/.venv/lib/python3.13/site-packages/IPython/core/pylabtools.py:158: UserWarning: Glyph 8722 (\N{MINUS SIGN}) missing from font(s) NanumSquareRound.
  fig.canvas.print_figure(bytes_io, **kw)

6. 짧게 학습해 보기 — 40스텝

실제 실험(10,000스텝, 75분)과 같은 옵티마이저·학습률 일정·손실이고 스텝 수만 짧습니다. 손실이 내려가고, 검증 이미지의 PSNR이 오르는지 봅니다.

def fg_psnr(pred8, hr8, mask):
    # 전경(alpha>0) 픽셀만의 8비트 RGB PSNR
    d = pred8.astype(np.float64)[mask] - hr8.astype(np.float64)[mask]; mse = float(np.mean(d**2))
    return 99.0 if mse == 0 else 10*np.log10(255**2/mse)

def eval_val(model, ds, idxs):
    # 검증 몇 장을 배치 없이 한 장씩 추론해 전경 PSNR을 잰다
    model.eval(); out = []
    with torch.no_grad():
        for i in idxs:
            it = ds[i]; p = model(it['lr'][None].cuda())[0]
            out.append(fg_psnr(to_img(p), to_img(it['hr']), it['mask'][0].numpy() > 0))
    return out

val_idx = [0, 40, 80, 120, 160, 200, 240]
before = eval_val(model, val_o, val_idx)
optimizer = torch.optim.Adam(model.parameters(), lr=config['learning_rate'], betas=config['betas'])
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, config['total_steps'], eta_min=config['min_learning_rate'])
history = []; t0 = time.perf_counter()
for step, batch in enumerate(loader(train_o, config), start=1):
    metrics = train_step(model, optimizer, batch, config, step-1); scheduler.step(); history.append(metrics)
    if step % 5 == 0 or step == 1:
        print(f"step {step:3d} · loss {metrics['loss']:.4f} · 전경 {metrics['foreground_l1']:.4f} · 배경 {metrics['background_l1']:.4f} · grad {metrics['gradient_norm']:.2f} · lr {optimizer.param_groups[0]['lr']:.2e}")
print(f"40스텝 {time.perf_counter()-t0:.0f}초 · GPU 최대 {torch.cuda.max_memory_allocated()/2**30:.1f} GiB")
after = eval_val(model, val_o, val_idx)
fig, ax = plt.subplots(1, 2, figsize=(11, 3.4))
ax[0].plot([h['loss'] for h in history], label='loss'); ax[0].plot([h['foreground_l1'] for h in history], label='전경 L1'); ax[0].plot([h['background_l1'] for h in history], label='배경 L1'); ax[0].set_xlabel('step'); ax[0].legend(); ax[0].set_title('학습 손실 (40스텝)')
ax[1].bar(np.arange(len(val_idx))-0.2, before, 0.4, label='학습 전'); ax[1].bar(np.arange(len(val_idx))+0.2, after, 0.4, label='40스텝 후'); ax[1].set_xticks(range(len(val_idx))); ax[1].set_xticklabels([val_o.rows[i]['name'][:6] for i in val_idx], fontsize=8); ax[1].set_ylabel('전경 PSNR (dB)'); ax[1].legend(); ax[1].set_title('검증 7장 PSNR')
plt.tight_layout(); plt.show()
print('검증 7장 평균 PSNR: 학습 전', f"{np.mean(before):.2f} dB → 40스텝 후 {np.mean(after):.2f} dB", '(실제 실험은 10,000스텝에서 12.2 → 16.8 dB)')

step   1 · loss 0.1422 · 전경 0.1400 · 배경 0.0214 · grad 0.87 · lr 2.00e-05

step   5 · loss 0.1114 · 전경 0.1094 · 배경 0.0199 · grad 0.53 · lr 1.93e-05

step  10 · loss 0.0833 · 전경 0.0809 · 배경 0.0236 · grad 0.52 · lr 1.74e-05

step  15 · loss 0.0995 · 전경 0.0974 · 배경 0.0218 · grad 0.36 · lr 1.44e-05

step  20 · loss 0.1114 · 전경 0.1088 · 배경 0.0262 · grad 0.46 · lr 1.10e-05

step  25 · loss 0.1101 · 전경 0.1074 · 배경 0.0267 · grad 0.47 · lr 7.56e-06

step  30 · loss 0.1201 · 전경 0.1179 · 배경 0.0223 · grad 0.49 · lr 4.64e-06

step  35 · loss 0.1347 · 전경 0.1326 · 배경 0.0215 · grad 0.53 · lr 2.69e-06

step  40 · loss 0.0880 · 전경 0.0853 · 배경 0.0269 · grad 0.35 · lr 2.00e-06
40스텝 20초 · GPU 최대 27.6 GiB

검증 7장 평균 PSNR: 학습 전 12.06 dB → 40스텝 후 13.57 dB (실제 실험은 10,000스텝에서 12.2 → 16.8 dB)

# 같은 검증 이미지 한 장: 입력 → 학습 전 출력은 위 5장에서 봤고, 여기서는 40스텝 후 출력과 정답을 나란히
i = val_idx[1]; it = val_o[i]
with torch.no_grad(): p = model(it['lr'][None].cuda())[0]
show([up(to_img(it['lr'])), to_img(p), to_img(it['hr'])], [f"입력 LR · {val_o.rows[i]['name']}", f'40스텝 후 출력 · {after[1]:.2f} dB', 'HR 정답'])
del model, optimizer; torch.cuda.empty_cache()

7. 실제 학습은 어떻게 돌렸나

위 코드에 스텝 수만 10,000으로 늘리고(total_steps=10000, workers=2), 500스텝마다 검증 42장을 채점해 최고 가중치를 고르고, 1,000스텝마다 가중치를 저장하며 전시 8종을 추론한 것이 experiments/swinir_point/run.py입니다.

cd ~/prj/maple/6_upscale && /home/son/prj/ComfyUI/.venv/bin/python -m experiments.swinir_point.run outline

• 결과 페이지: /maple_upscale/swinir/outline/ · 결과 갤러리: /maple_upscale/result_gallery/
• 설정: 배치 8 · Adam(0.9, 0.99) · 학습률 2e-5 → 2e-6 cosine · grad clip 1.0 · FP32(TF32 끔) · HR 192 패치 · 시드 20260912