模仿Keras的进度条

挺喜欢 Keras 在模型训练时显示的那个进度条, 打算移植到 PyTorch 的训练代码上.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import sys
import time

def print_msg(message, line_break=False):
    """打印消息到标准输出"""
    if line_break:
        sys.stdout.write("\r" + message + "\n")
    else:
        sys.stdout.write("\r" + message)
    sys.stdout.flush()
    
def time_duration_value():
    return "6s 12ms/step"

def loss_value():
    return "loss: 0.0377"

def accuracy_value():
    return "accuracy: 0.9890"

epoch_count = 5
for i in range(epoch_count):
    print_msg(f"Epoch {i+1:{len(str(epoch_count))}}/{epoch_count}", line_break=True)
    batch_count = 10
    for j in range(batch_count):
        message = f"{j+1:{len(str(batch_count))}}/{batch_count} [{'=' * (j+1):{batch_count}}]"
        message += " - " + time_duration_value()
        message += " - " + loss_value()
        message += " - " + accuracy_value()
        print_msg(message)
        time.sleep(0.1)
    print_msg(message, line_break=True)

效果:

1
2
3
4
5
6
Epoch 1/5
10/10 [==========] - 6s 12ms/step - loss: 0.0377 - accuracy: 0.9890
Epoch 2/5
10/10 [==========] - 6s 12ms/step - loss: 0.0377 - accuracy: 0.9890
Epoch 3/5
 5/10 [=====     ] - 6s 12ms/step - loss: 0.0377 - accuracy: 0.9890