[ultralytics] 행렬곱 두 번으로 bilinear 보간 대체하기: deterministic 모드에서 Ultralytics semantic 학습 최대 3.9배 가속

PR 링크: ultralytics/ultralytics#27154 상태: Merged | 변경: +37 / -10

들어가며

Ultralytics 학습은 기본값이 deterministic=True입니다. 이 설정에서는 BaseTrainer가 init_seeds()를 통해 torch.use_deterministic_algorithms(True)를 켭니다. 그런데 CUDA 같은 non-CPU 디바이스에서 F.interpolate(mode="bilinear")의 backward(upsample_bilinear2d_backward)는 atomic add 기반이라 결정적이지 않습니다. 그래서 PyTorch는 deterministic 모드에서 이 연산을 Python 레벨 decomposition으로 바꿔 실행하고, 이 경로는 상당히 느립니다.

semantic segmentation, depth estimation, segment proto loss는 모두 매 iteration마다 bilinear resize를 여러 번 호출합니다. 따라서 이 decomposition이 그대로 학습 시간의 병목이 됩니다. 이 PR은 bilinear 보간을 두 번의 matmul로 다시 표현해서 이 병목을 없앱니다.

핵심 아이디어: bilinear는 separable하다

bilinear 보간은 H축 linear 보간과 W축 linear 보간을 차례로 적용한 것과 같습니다. 각 축의 보간은 입력에 대한 선형 변환이므로 행렬로 쓸 수 있습니다.

out = Wh @ x @ Ww^T    # Wh: (out_h, in_h), Ww: (out_w, in_w)

이 경로에서는 matmul의 backward 역시 matmul입니다. 그래서 decomposition 없이 결정적이고 빠르게 실행됩니다.

코드 분석

ultralytics/utils/torch_utils.py: 새 헬퍼 interpolate_bilinear

def interpolate_bilinear(x, size=None, scale_factor=None, align_corners=False):
    if x.device.type == "cpu" or not torch.are_deterministic_algorithms_enabled():
        return F.interpolate(x, size, scale_factor, "bilinear", align_corners)
    with autocast(False, x.device.type):  # keep the input dtype, as the native kernel does
        h, w = (  # per-axis weights from the native 1D linear kernel, which deterministic mode keeps native
            F.interpolate(torch.eye(n, device=x.device, dtype=x.dtype)[None], s, scale_factor, "linear", align_corners)
            for n, s in zip(x.shape[-2:], size or (None, None))
        )
        return h[0].T @ x @ w[0]

눈여겨볼 지점은 네 가지입니다.

1. 가중치 행렬을 직접 계산하지 않습니다. torch.eye(n)[None]은 shape이 (1, n, n)인 텐서입니다. 이를 (N=1, C=n, L=n)으로 보고 1D "linear" 보간을 적용하면, 결과 (1, n, s)의 각 행이 "입력 픽셀 i가 출력 위치들에 기여하는 가중치"가 됩니다. 덕분에 align_corners와 scale_factor 처리 같은 좌표 계산 규칙을 PyTorch 네이티브 커널에 그대로 위임합니다. 직접 수식을 짜면 생길 수 있는 off-by-half-pixel 버그를 원천적으로 피하는 셈입니다. 게다가 1D linear 커널은 deterministic 모드에서도 네이티브로 유지되고, 상수 행렬에 대한 forward라서 backward도 필요 없습니다.
2. 브로드캐스팅 matmul: h[0].T는 (out_h, H), x는 (B, C, H, W), w[0]는 (W, out_w)입니다. 따라서 h[0].T @ x는 (B, C, out_h, W)가 되고, 여기에 @ w[0]를 하면 (B, C, out_h, out_w)가 됩니다. reshape 없이 한 줄로 끝납니다.
3. autocast(False, ...): AMP 환경에서 matmul은 autocast에 의해 fp16으로 캐스팅됩니다. 하지만 네이티브 interpolate는 입력 dtype을 유지합니다. 이 블록이 동작을 동일하게 맞춰 주기 때문에, 호출부의 dtype 가정이 깨지지 않습니다.
4. fallback 조건: CPU이거나 deterministic 모드가 꺼져 있으면 기존 F.interpolate를 그대로 호출합니다. 따라서 CPU 학습, export, deterministic=False 실행은 전혀 바뀌지 않습니다.

ultralytics/utils/loss.py: loss 경로 교체

semantic loss는 main logits와 aux logits를 모두 마스크 크기로 upsample합니다.

# Before
preds = F.interpolate(preds, size=masks.shape[1:], mode="bilinear", align_corners=False)
# After
preds = interpolate_bilinear(preds, masks.shape[1:])

segment proto loss와 depth loss도 같은 방식으로 바뀌었습니다. depth는 원래 의미대로 align_corners=True를 유지합니다.

# Before
pred_depth = F.interpolate(pred_depth, size=gt_depth.shape[-2:], mode="bilinear", align_corners=True)
# After
pred_depth = interpolate_bilinear(pred_depth, gt_depth.shape[-2:], align_corners=True)

semantic 학습이 3.9배로 가장 크게 빨라진 이유가 여기에 있습니다. 출력 해상도가 크고, 이 경로를 backward까지 거치기 때문입니다.

ultralytics/nn/modules/head.py: depth head의 FPN 업샘플

# Before
out = F.interpolate(out, scale_factor=2, mode="bilinear", align_corners=True)
# After
out = interpolate_bilinear(out, scale_factor=2, align_corners=True)

이 호출은 size 대신 scale_factor를 씁니다. 그래서 헬퍼의 size or (None, None) 분기가 필요합니다. 주석에 적힌 대로 align_corners=True는 공개 가중치에 박혀 있는 값이라 반드시 그대로 보존해야 합니다.

validator들 (semantic/val.py, depth/val.py)

검증 단계의 resize도 같은 헬퍼로 교체했습니다. 그 결과 depth/val.py에서는 torch.nn.functional import 자체가 사라졌습니다.

리뷰에서 나온 논의

머지 직전 메인테이너 glenn-jocher가 두 가지를 수정했습니다.

• 디바이스 조건 단순화: 원래 코드는 TORCH_2_9와 is_cuda를 조합해 검사했습니다. 이 방식은 torch 2.5~2.8의 XPU를 놓쳤습니다. 그래서 "CPU가 아니고 deterministic이 켜져 있으면"이라는 조건으로 바꿨는데, 이것이 PyTorch가 실제로 decomposition하는 조건과 일치합니다. 또 Tensor.is_cpu는 프로젝트의 최소 지원 버전인 torch 1.8에 없어서 x.device.type == "cpu"를 사용했습니다.
• tracing guard 제거: 자동 리뷰 봇은 deterministic 모드에서 ONNX dynamic=True로 export하면 torch.eye(n)이 trace 시점의 크기로 고정될 수 있다고 지적했습니다. 이에 대한 반론은 이렇습니다. deterministic 모드는 BaseTrainer 안에서만 켜지고 finally의 unset_deterministic()으로 꺼지므로, export 중에는 이 경로에 도달할 수 없습니다. 게다가 torch.jit.is_tracing()은 dynamo나 torch.export를 잡아내지도 못합니다. 도달 불가능한 상태를 위한 방어 코드는 복잡도만 늘린다는 판단입니다.

왜 이게 좋은가

학습 (deterministic=True)	before	after	speedup
semantic nc19 b8, 3 epochs, T4	626 s	159 s	3.9x
depth nyu b16, 5 epochs, T4	189 s	143 s	1.3x
semantic nc19 b4, 1 epoch, M4 MPS	80.7 s	45.8 s	1.8x

개별 호출 지점만 놓고 보면 5~30배 빨라졌습니다. 정확도 측면에서 forward 오차는 9.5e-7 이하, gradient 오차는 9.5e-5 이하로 fp32 반올림 수준입니다. 반복 실행에서는 bitwise로 동일한 metric이 나와 결정성도 유지됩니다. 50 epoch 비교에서도 mIoU와 delta1은 seed 간 편차 범위 안에 있었습니다.

여기서 얻을 수 있는 일반적인 교훈은 다음과 같습니다.

1. 결정성 플래그는 공짜가 아닙니다. use_deterministic_algorithms는 일부 연산을 느린 대체 경로로 바꿉니다. 프로파일링으로 이런 경로를 찾아보면 의외로 큰 이득이 숨어 있습니다.
2. 선형 연산은 행렬로 바꿀 수 있습니다. resize, pooling 같은 선형 변환을 matmul로 표현하면 GPU에서 가장 최적화된 커널(GEMM)을 쓰게 되고, backward도 자동으로 결정적이 됩니다.
3. 수식을 재구현하지 말고 네이티브 커널에서 가중치를 뽑아 쓰세요. identity 행렬을 1D 커널에 통과시키는 트릭은 정합성을 공짜로 보장해 줍니다.
4. fallback으로 영향 범위를 최소화하세요. 최적화는 문제가 실제로 발생하는 조건에서만 활성화하고, 나머지 경로는 바이트 단위로 동일하게 유지하는 것이 안전한 변경의 핵심입니다.

참고 자료

• torch.nn.functional.interpolate — PR에서 대체·활용한 보간 함수
• torch.use_deterministic_algorithms — deterministic 모드와 영향을 받는 연산 목록
• torch.are_deterministic_algorithms_enabled — 헬퍼의 분기 조건에 사용된 API
• torch.eye — 보간 가중치 행렬 생성에 사용
• torch.matmul — 브로드캐스팅 행렬곱
• Automatic Mixed Precision (torch.amp) — autocast 비활성화로 dtype 유지
• Reproducibility — PyTorch 재현성 노트

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.