[PyTorch] 모델 efficiency 측정하기 (used gpu memory / parameter 개수 / Inference time)

 

 

모델 efficiency를 측정하다가, 아래와 같은 수치가 필요해져서 한번 정리해놓는다. 

 

# GPU memory consumption

참고 doc

def to_MB(a):
    return a/1024.0/1024.0
    
print(f"After model to device: {to_MB(torch.cuda.memory_allocated()):.2f}MB")

위에서 출력된 nividia-smi 와 다르니 주의해아한다. 왜냐면 nvidia-smi는 reserved memory를 반환하고, 위 코드는 allocated memory를 반환하기 때문이다. 흔히 reserved memory는 캐시까지 포함한 메모리를 일컫는다고 한다. 

 

# Parameter count 

이건 단순 trainable parameter의 개수를 count 해주는 코드이다.

model_parameters = filter(lambda p: p.requires_grad, model.parameters())
params = sum([np.prod(p.size()) for p in model_parameters])

 

# Training/Inference time

아래와 같이 python의 time 모듈로 측정하지 않도록 주의하자. 

 s = time.time()
 _ = model(dummy_input)
curr_time = (time.time()-s )*1000

그 대신 아래와 같이 torch.cuda.Event를 이용하면 된다.

starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)   # Time evaluation 
starter.record()

# 측정하고 싶은 모델 프로세스  

ender.record()
torch.cuda.synchronize()
curr_time = starter.elapsed_time(ender)
print(curr_time)

이 글을 보는 사람들의 연구와 paper 작성을 응원합니다 !!! 🔥