[Python] tqdm 이용해 모델 train 시 progress bar 표시하기

tqdm module을 이용하면 iteration 진행 상황을 progress bar로 나타낼 수 있다.

간단한 예제는 다음과 같다.

 

from tqdm import tqdm
from time import sleep

for i in tqdm(range(100)):
  sleep(0.1)

100%|██████████| 100/100 [00:10<00:00,  9.76it/s] # output

 

다음과 같이 itreration 부분을 tqdm으로 감싸주면 iteration이 진행되는 progress bar를 확인할 수 있다.

이를 딥러닝 model train 시 epoch 당 progress를 확인하기 위해 적용해 보자.

 

for epoch in range(num_epoch):

    with tqdm(train_loader, unit="batch") as tepoch:
        
        for x, y in tepoch:
            
            tepoch.set_description(f"Epoch {epoch+1}")

            optimizer.zero_grad()
                
            output = model(x)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()
            
            tepoch.set_postfix(loss=loss.item())

Epoch 1: 100%|██████████| 3/3 [00:00<00:00,  6.01batch/s, loss=0.82]

우선 train dataloader를 tqdm을 이용해 감싸준다. dataloader가 batch 단위로 데이터를 뱉어낼 것이기 때문에 Unit 이름은 batch로 설정하자. 이렇게 하면 초당 몇 개의 batch가 모델로 들어가는지 확인할 수 있다.

 

set_description을 이용해 iteration 이름을 정의할 수 있다.

현재 Epoch 당 progress를 확인하고 있기 때문에 epoch으로 설정해주었다.

 

set_postfix를 이용하면 원하는 변수의 update 상황에 대해 확인할 수 있다.

 

---

Reference : https://towardsdatascience.com/training-models-with-a-progress-a-bar-2b664de3e13e