python在类中手动定义标准化输出函数

news2024/9/21 12:48:46

在深度学习等训练框架中，有时候需要对模型的名称、参数量、训练进度、中间结果等进行标准化输出，从而方便实时查看代码运行情况，这时，可以在类中手动定义如下标准化的输出函数，然后在需要输出的地方进行调用即可。

首先，先介绍和定义一个对数字（训练步数）进行标准化显示的函数，如：1->1.0; 1000->1.0K; 10000->10.0K; 10000000->10.0M

def human_format(num):
    magnitude = 0
    while num >= 1000:
        magnitude += 1
        num /= 1000.0
    # add more suffixes if you need them
    return '{:3.1f}{}'.format(num, [' ', 'K', 'M', 'G', 'T', 'P'][magnitude])

然后，可以定义如下两个标准化输出的函数，其中 verbose() 用于静态显示要输出的内容，并进行左对齐；progress() 用于动态显示实时更新的内容

import sys

Class XXX():
    def __init__(self, ...):
        ...

    def verbose(self, msg):
        ''' Verbose function for print information to stdout'''
        if self.paras.verbose:
            if type(msg) == list:
                for m in msg:
                    print('[INFO]', m.ljust(100)) # 左对齐，在当前字符串的右侧填充空格，使整个字符串长度达到100个字符
            else:
                print('[INFO]', msg.ljust(100))

    def progress(self, msg):
        ''' Verbose function for updating progress on stdout (do not include newline) '''
        if self.paras.verbose:
            sys.stdout.write("\033[K")  # Clear line 这行代码使用sys.stdout.write方法向标准输出写入特殊的ANSI转义码\033[K，用于清除当前行的内容
            print('[{}] {}'.format(human_format(self.step), msg), end='\r') # end='\r'参数告诉print函数在打印完消息后不换行，而是回到行首，以便在下一次调用progress方法时覆盖前一次的输出，从而实现进度更新的效果。

例如，在如下调用时，会标准化输出如下：

self.verbose("Saved checkpoint (step = {}, {} = {:.2f}) and status @ {}".
            format(human_format(self.step), metric, score, ckpt_path))

[INFO] Saved checkpoint (step = 12.0K, wer = 0.62) and status @ ckpt/asr_hybrid_english_sd0/best_att.pth

if (self.step == 1) or (self.step % self.PROGRESS_STEP == 0):
    self.progress('Tr stat | Loss - {:.2f} | Grad. Norm - {:.2f} | {}'
        .format(total_loss.cpu().item(), grad_norm, self.timer.show()))