Tensorboard 训练分类算法的tensorboard可视化示例

Posted 2022-11-22 洪流之源

该示例实现了tensorboard对训练数据图像、模型计算图、学习率、损失值、准确率、梯度、模型权值参数的可视化：

# 监控loss, accuracy, weights, gradients
import numpy as np
import math
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets
from torchvision import models
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

if __name__ == \'__main__\':

    # 参数设置
    NUM_CLASSES = 100
    MAX_EPOCH = 100
    BATCH_SIZE = 64
    INIT_LR = 0.01
    LRF = 0.1 # 倍率因子，最后学习率会降到初始学习率的0.1倍
    LOG_INTERVAL = 10
    VAL_INTERVAL = 1
    SAVE_CHECKPOINTS_DIR = \'checkpoints\'

    # 构建 SummaryWriter
    summary_writer = SummaryWriter(log_dir=\'logs\')

    if not os.path.exists(SAVE_CHECKPOINTS_DIR):
        os.makedirs(SAVE_CHECKPOINTS_DIR)
    
    device = torch.device("cuda:0" if torc