📈 FLOPs 카운팅 메커니즘의 구현
실험을 몇 번 반복하다 보면, 어느 순간부터 "정확도는 잘 나오는데 왜 이렇게 느리지?" 같은 질문이 계속 남는다. 특히 Lucid처럼 CPU/Metal(GPU) 백엔드, lazy materialization(loss.eval()), 그리고 자체 연산/모듈 스택을 갖춘 프레임워크에서는, 단순한 wall-clock time 비교가 원인을 깔끔히 설명해주지 못한다.
그래서 이번 개발 일지에서는 Lucid에 FLOPs(정확히는 MACs 기준) 카운팅 메커니즘을 도입한 과정을 정리한다. 왜 필요하다고 느꼈는지, 어떤 자료구조를 차용했는지, 그리고 실제로 이를 코드로 실체화하기 위해 어떤 구현을 했는지까지 "프레임워크 내부 설계" 관점으로 풀어보는 것이 목표다.
🧭 FLOPs 카운팅의 필요성
딥러닝 프레임워크를 만들면서 "속도"는 항상 중요한 지표인데, 속도를 직접 재는 방식에는 크게 두 가지 문제가 있다.
- 환경 의존성: CPU/Metal(GPU) 설정, 커널 최적화 수준, 캐시/메모리 상태가 시간에 크게 영향을 준다.
- 원인 분리가 어렵다: 느린 원인이 "연산량(FLOPs)이 많아서"인지, "메모리 이동/형 변환/그래프 평가 타이밍" 때문인지 구분이 어렵다.
이때 FLOPs는 wall-clock time을 대체하지는 않지만, 최소한 다음 질문에 답할 수 있게 해준다.
- "이 모델/이 forward 패스가 계산적으로 얼마나 비싼가?"
- "같은 모델인데 결과가 다르면 실제 연산 그래프가 달라진 건가?"
- "최적화(예: 연산 구현 변경, fused path 도입)의 효과가 연산량 vs. 실행시간 중 어디서 오는가?"
즉 FLOPs는 성능 분석의 좌표계 역할을 한다.
🧰 설계 목표
Lucid의 FLOPs 카운팅은 아래 목표를 기준으로 설계했다.
- Opt-in: 기본 실행에서는 오버헤드가 없어야 한다.
- 그래프 재사용: 이미 존재하는 autograd 그래프(
Tensor._prev,Tensor._op)를 그대로 활용한다. - 연산 단위 확장성: 연산마다 FLOPs 정의가 다르므로, op별로 쉽게 커스터마이즈 가능해야 한다.
- CPU/GPU 공통: 어떤 디바이스 경로로 실행되든 동일한 논리로 FLOPs를 계산한다(연산 정의는 동일, 구현만 다름).
🧷 핵심 아이디어
Lucid는 각 연산을 operation 객체로 캡슐화하고(lucid/_backend/core.py), 실제 텐서 연산 호출은 func_op 데코레이터가 감싼 wrapper에서 처리된다. 이 구조가 FLOPs 카운팅에 매우 잘 맞는다.
1️⃣ FLOPs 카운팅 on/off 스위치
전역 플래그를 두고, context manager로만 켜지게 했다(lucid/__init__.py).
_flops_enabled: bool = False
@contextmanager
def count_flops():
global _flops_enabled
prev_state = _flops_enabled
_flops_enabled = True
try:
yield
finally:
_flops_enabled = prev_state이 방식의 장점은 명확하다. "평소에는 꺼져 있다가", 필요할 때만 계산 그래프를 보존하고 FLOPs를 기록한다.
2️⃣ 연산이 실행될 때 해당 op의 FLOPs를 계산해 저장
func_op wrapper에서 FLOPs 카운팅이 켜져 있으면, 각 operation 인스턴스에 op_self.flops = op_self.__flops__(...)를 기록한다(lucid/_backend/core.py).
flops_enabled = lucid.flops_enabled()
track_graph = flops_enabled or (grad_enabled and requires_grad)
if flops_enabled:
op_self.flops = op_self.__flops__(*new_args, **kwargs)여기서 핵심은 track_graph 조건이다. gradient가 필요 없는 inference라도 FLOPs 카운팅을 하려면 그래프를 추적해야 하므로, flops_enabled가 켜진 경우에도 _op/_prev를 유지하도록 설계했다.
3️⃣ 최종 output Tensor에서 그래프를 순회하며 FLOPs 합산
각 Tensor는 자신을 만든 op(_op)와 입력 텐서들(_prev)을 알고 있다. 따라서 output에서 시작해 그래프를 거슬러 올라가며 op.flops를 더하면 된다(lucid/_tensor/tensor.py).
@property
def flops(self) -> int:
total = 0
queue = deque([self])
visited = set()
while queue:
cur = queue.popleft()
visited.add(cur)
if cur._op is not None:
total += cur._op.flops
for nxt in cur._prev:
if nxt not in visited:
queue.appendleft(nxt)
return total이 구현은 "그래프 순회"라는 측면에서 DFS에서 영감을 받았다. 핵심은 output 텐서에서 시작해 dependency를 따라 아래로 내려가며(DAG를 거슬러) 모든 노드를 한 번씩 방문한다는 점이다.
Lucid의 autograd 그래프는 트리(tree)가 아니라 보통 DAG다. 예를 들어 y = a + a 같은 경우, 동일한 텐서 a가 두 경로로 참조된다. 이때 단순 재귀 DFS로 "부모 → 자식"을 계속 따라가면, 같은 노드를 여러 번 방문할 수 있고 FLOPs가 중복 합산될 수 있다. 그래서 visited를 두고 "이미 방문한 텐서는 스킵"하는 규칙을 넣었다. 이 규칙이 들어가면 그래프 구조가 어떻든 각 텐서(노드)당 최대 1회만 처리하게 되고, 시간 복잡도는 사실상 로 안정된다.
또 하나의 포인트는 "재귀 DFS"가 아니라 반복(iterative) DFS 스타일로 작성했다는 점이다. 파이썬 재귀는 깊은 그래프에서 recursion limit에 걸릴 수 있고, 프레임워크 내부에서 예외가 섞이면 디버깅도 까다로워진다. 그래서 deque를 스택처럼 사용했다.
popleft()로 현재 노드를 꺼내고- 자식(여기서는
_prev)을appendleft()로 다시 왼쪽에 넣는다
이 패턴은 결과적으로 LIFO에 가까운 "스택 기반 DFS" 순회가 된다. (만약 append()를 썼다면 FIFO에 가까운 BFS 스타일이 되지만, FLOPs 합산은 방문 순서에 의존하지 않기 때문에 둘 다 가능하다. 다만 DFS 쪽이 구현 감각상 "그래프를 한 갈래로 깊게 따라간다"는 직관이 더 강했다.)
마지막으로 cur._op가 있으면 cur._op.flops를 더하는데, 이때 "op"는 "결과 텐서를 만든 연산"이므로, output에서 시작해 입력으로 내려가며 방문하면 실제로 실행된 연산들의 FLOPs를 누적하게 된다. 이 구조 덕분에 FLOPs 카운팅은 별도의 트레이서 없이도, 기존 autograd 그래프를 그대로 재활용해 구현될 수 있었다.
⏳ 연산별 FLOPs 정의
operation 베이스는 기본적으로 __flops__를 0으로 둔다(lucid/_backend/core.py). 즉, FLOPs 정의가 필요한 연산은 각각의 클래스에서 오버라이드만 해주면 된다.
🔹 원소 단위(elementwise) 연산
예를 들어 add는 broadcasting 이후 output shape의 element 수를 FLOPs로 둔다(lucid/_func/bfunc.py).
def __flops__(self, a: Tensor, b: Tensor) -> int:
out_shape = np.broadcast_shapes(a.shape, b.shape)
return int(np.prod(out_shape))🔹 행렬곱(matmul)
Lucid의 matmul은 배치까지 고려해 batch * m * n * k 형태로 MACs를 계산한다(lucid/_func/bfunc.py).
def __flops__(self, a: Tensor, b: Tensor) -> int:
# ...
return batch_size * m * n * k여기서 "FLOPs"는 흔히 말하는 2×MACs(곱+더하기) 정의가 아니라, MACs 기반으로 잡았다. (그래야 matmul/conv 같은 핵심 연산들이 같은 기준으로 비교된다.)
🔹 합성곱(conv)
Conv는 모델 FLOPs의 상당 부분을 차지하므로, 별도로 conv backend op(conv_nd)에 __flops__를 추가했다(lucid/_backend/conv.py).
- weight shape:
(C_out, C_in/groups, *kernel) - output shape:
(N, C_out, *out_dims) - MACs per output element:
(C_in/groups) * prod(kernel)
즉,
bias는 conv(...)에서 output에 더해지는 elementwise add로 처리되므로(별도의 add op), 자연스럽게 FLOPs에 반영된다.
🧪 실사용 예제
1️⃣ 단순 forward의 FLOPs 얻기
import lucid
import lucid.nn as nn
model = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(16, 32, kernel_size=3, padding=1),
)
x = lucid.random.rand((1, 3, 32, 32))
with lucid.count_flops():
y = model(x)
print(f"FLOPs(MACs): {y.flops:,}")여기서 중요한 포인트는 with lucid.count_flops(): 블록이 있어야 y.flops가 의미를 갖는다는 점이다. (블록 밖에서는 그래프가 추적되지 않을 수 있다.)
2️⃣ 모델 요약 출력에서 FLOPs 확인
Lucid는 summarize에서 내부적으로 count_flops()를 사용해 총 FLOPs를 출력한다(lucid/models/util.py).
from lucid.models.util import summarize
from lucid.models.imgclf.resnet import resnet_18
model = resnet_18(num_classes=10)
summarize(model, input_shape=(1, 3, 32, 32))요약 테이블 하단에 Total FLOPs가 함께 출력되며, 이는 "이 모델이 대략 어느 정도 계산량을 요구하는가"를 빠르게 판단하는 기준선이 된다.
🧯 한계와 다음 단계
FLOPs는 어디까지나 "계산량"의 지표다. 따라서 다음 한계가 있다.
- 메모리/IO 비용 미반영: 실제 속도는 메모리 대역폭/캐시/디바이스 이동에 좌우될 수 있다.
- 정의의 차이: FLOPs를 2×MACs로 볼지, MACs로 볼지는 커뮤니티마다 다르다. Lucid는 현재 구현(특히 matmul)과의 일관성을 위해 MACs 기준을 채택했다.
- 커버리지: 모든 연산이
__flops__를 갖고 있지는 않다. 핵심 경로부터 채우되, 사용 빈도/중요도에 따라 점진적으로 확장하는 편이 현실적이다.
다음 단계로는
- 주요 nn.functional 경로(특히 attention 계열) FLOPs 정의 확장,
- forward hook을 활용한 모듈 단위 FLOPs breakdown(레이어별 표),
- FLOPs와 함께 메모리 footprint/activation size까지 같이 보여주는 "프로파일 요약"
같은 방향으로 발전시키는 것이 자연스럽다.
✅ 정리
이번 작업의 핵심은 "새로운 프로파일러를 따로 만들지 않고", Lucid가 이미 갖고 있던 연산 그래프 구조(operation + Tensor DAG) 위에 FLOPs라는 메타데이터를 얹었다는 점이다. count_flops()로 오버헤드를 통제하고, 각 op가 __flops__만 제공하면 전체 그래프 FLOPs를 자동으로 합산할 수 있도록 만들면서, 성능 분석과 구현 개선을 위한 공통 언어(연산량)를 Lucid 내부에 마련했다.
