[Pytorch專欄]使用tensorboardX進行訓練可視化

訓練時候可視化loss曲線非常有用，可以很好的觀察是否過擬合，還是存在欠擬合，還可以直接觀察測試精度

感謝作者開源：下載鏈接，直接使用pip安裝也是可以的：pip install tensorboardX, tensorflow-gpu

注意安裝的時候必須連帶tensorflow一起安裝，此時會將tensorboard安裝好，否則會提示：tensorboard: command not found。tensorboardX藉助於tensorboard進行可視化

 

今天簡單講講loss曲線的可視化，以maskrcnn-benchmark爲例。首先引入一個TensorboardXWriter

from tensorboardX import SummaryWriter


class TensorboardXWriter(object):
    """Writes for writing training info to file"""

    @staticmethod
    def init(log_dir):
        """
        Args:
            log_dir (str): path to log directory 
        """
        TensorboardXWriter.writer = SummaryWriter(log_dir=log_dir)

    @staticmethod
    def write_scalar(names, values, iter):
        #注意傳入names和values都是list
        #names表示標題名稱，values表示Y軸數值，iter表示迭代次數X軸
        for var_name, var_value in zip(names, values):
            TensorboardXWriter.writer.add_scalar(var_name, var_value, iter)

    @staticmethod
    def write_scalars(names, values, iter):
        #注意傳入names和values都是list
        #names表示標題名稱，values表示Y軸數值，iter表示迭代次數X軸
        for var_name, var_value in zip(names, values):
            TensorboardXWriter.writer.add_scalars(var_name, var_value, iter)
    
    @staticmethod
    def close():
        TensorboardXWriter.writer.close()

這個類包含三個函數，一個是init函數，用於初始化tb_event目錄，還有兩個函數write_scalar和write_scalars，分別用於加載一個數值和一個列表，但是注意必須封裝成list。

value參數如果是PyTorch scalar tensor，則需要調用 .item() 方法獲取其數值。一般會使用 add_scalar 方法來記錄訓練過程的 loss、accuracy、learning rate 等數值的變化

在tools/train_net.py中初始化tb_event

tf_event_dir = os.path.join(output_dir, "tb_events")
try:
    os.makedirs(tf_event_dir)
except FileExistsError:
    pass

TensorboardXWriter.init(log_dir=tf_event_dir)

這樣完成了初始化操作，然後在maskrcnn_benchmark/engine/trainer.py中將loss進行可視化

if is_main_process():
    TensorboardXWriter.write_scalars(
        ["train/losses", "train/time"],
        [
            meters.get_metric(
                 metric_names=[
                    "loss",
                    "loss_box_reg",
                    "loss_classifier",
                    "loss_objectness",
                    "loss_rpn_box_reg"
                ]
            ),
            meters.get_metric(metric_names=["time", "data"]),
        ],
        iteration,
    )    

還需要給meters添加get_metric方法，這個比較簡單，關於庫函數的導入這裏也不羅嗦了，需要特別提醒的一點是必須要進行is_main_process的判斷，否則在多卡訓練的時候會出現KeyError或者FileExistError創建目錄失敗。這樣使用靜態函數進行封裝的好處是無需進行對象的傳遞。後臺啓動tensorboard執行下述語句，--bind_all可以登錄查看服務器上啓動的tensorboard，--port 可以自己指定端口號，默認的是6006

nohup tensorboard --logdir ./ --bind_all > tb.log 2>&1 &

這樣就完成了loss曲線的可視化

使用tensorboardX還可以可視化網絡結構或者運行圖add_graph、直方圖add_histogram、 圖像add_image和嵌入向量add_embedding，可以參考這篇文章瞭解詳解PyTorch項目使用TensorboardX進行訓練可視化

固定網絡結構使用tensorboardX繪圖不如使用netron更加方便，或者使用hiddenlayer更加方便，可惜他們對兩階段的網絡支持都不好。如果讀者發現好的兩階段網絡可視化麻煩留言告知呀