使用TensorBoard可视化pytorch模型、数据以及模型训练过程

tensorboard常用函数列表如下：

torch.utils.tensorboard.writer.SummaryWriter	tensorboard的入口，创建event files
add_scalar(tag, scalar_value, global_step=None, walltime=None)	添加标量到event file
add_scalars(main_tag,tag_scalar_dict, global_step=None, walltime=None)	添加一组标量到event file
add_histogram(tag, values, global_step=None, bins='tensorflow', walltime=None, max_bins=None)	添加柱状图
add_image(tag, img_tensor, global_step=None, walltime=None, dataformats='CHW')	添加图片
add_images(tag, img_tensor, global_step=None, walltime=None, dataformats='NCHW')	添加图片
add_figure(tag,figure, global_step=None, close=True, walltime=None)	将图片经过matplotlib处理后添加到TensorBoard
add_video(tag, vid_tensor, global_step=None, fps=4, walltime=None)	添加视频
add_audio(tag, snd_tensor, global_step=None, sample_rate=44100, walltime=None)	添加音频
add_text(tag, text_string, global_step=None, walltime=None)	添加文本
add_graph(model, input_to_model=None, verbose=False)	添加网络结构图
add_embedding(mat, metadata=None, label_img=None, global_step=None, tag='default', metadata_header=None)	将高维数据映射到低维进行表示
add_pr_curve(tag, labels, predictions, global_step=None, num_thresholds=127, weights=None, walltime=None)	添加PR曲线
add_custom_scalars(layout)
add_mesh(tag, vertices, colors=None, faces=None, config_dict=None, global_step=None, walltime=None)
add_hparams(hparam_dict=None, metric_dict=None)	超参数
close()

下面通过具体代码实例演示pytorch模型中使用tensorboard模块进行模型、数据、训练指标的可视化：

1、加载数据集

%matplotlib inline
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.optim as optim
import torch.nn.functional as F
from matplotlib import pyplot as plt
import torchvision
import os
import sys

# 设置下载的pretrained模型保存位置
os.environ["TORCH_HOME"] = "./models"

# transforms，预处理
transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize((0.5, ), (0.5, ))
    ]
)

# datasets
trainset = torchvision.datasets.FashionMNIST("./data", download=True, train=True, transform=transform)
testset = torchvision.datasets.FashionMNIST("./data", download=True, train=False, transform=transform)

# dataloaders
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2, drop_last=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2, drop_last=False)

# constant for classes
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')

# helper function to show an image
# (used in the `plot_classes_preds` function below)
def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    if one_channel:
        plt.imshow(npimg, cmap="Greys")
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

for batch_images, batch_labels in trainloader:
    print(batch_images.shape)
    print(batch_labels.shape)
    for img in batch_images:
        print(img.shape)
        matplotlib_imshow(img, one_channel=True)
    break

2、定义网络模型

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()
net

3、定义损失函数和优化方法

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

4、tensorboard初始化

启动tensorboard命令: tensorboard --logdir=logs

访问tensorboard页面：http://localhost:6006

from torch.utils.tensorboard import SummaryWriter

# 创建TensorBoard写入对象
writer = SummaryWriter("./runs/fashion_mnist_experiment_1")

# 将训练数据图像写入TensorBoard
dataiter = iter(trainloader)
images, labels = dataiter.next()

# 创建图像网格区域
img_grid = torchvision.utils.make_grid(images)
print(type(img_grid))

matplotlib_imshow(img_grid, one_channel=True)

5、add_image添加图片到tensorboard

writer.add_image("fasion mnist images", img_grid)

IMAGES标签下可以看到如下图像：

6、add_graph可视化网络模型结构

writer.add_graph(net, images)

GRAPHS标签下可以看到如下网络结构：

7、add_scalar添加标量，标量值随着global_step变化，就构成了变化趋势曲线

for i in range(10):
    writer.add_scalar("y = x^2", i ** 2, i)

SCALARS标签下可以看到如下变化曲线：

8、add_embedding将高维数据映射到低维进行表示

def select_n_random(data, labels, n=100):
    assert len(data) == len(labels)
    permutation_index = torch.randperm(len(data))
    return data[permutation_index][:n], labels[permutation_index][:n]

images, labels = select_n_random(trainset.data, trainset.targets)
class_labels = [classes[l] for l in labels]

print(images.shape, labels.shape)
print(images.unsqueeze(1).shape)
print(class_labels[:10])

features = images.view(-1, 28 * 28)
print(features.shape)

# features的形状是 (N, D）
# metadata对应类别的字符串
# label_img的形状是(N, C, H, W),也就是一个batch的图像数据
# D = C * H * W
writer.add_embedding(features, metadata=class_labels, label_img=images.unsqueeze(1))

PROJECTORS标签下可以看到如下三维图像：

9、add_figure、add_scalar、add_pr_curve记录模型训练准确率、损失函数、PR曲线、ROC曲线等

from sklearn.metrics import roc_curve
from sklearn.metrics import auc

# helper functions

# 计算输入图像的预测类别和对应的预测概率
def images_to_probs(net, images):
    '''
    Generates predictions and corresponding probabilities from a trained
    network and a list of images
    '''
    output = net(images)
    # convert output probabilities to predicted class
    _, preds_tensor = torch.max(output, 1)
    preds = np.squeeze(preds_tensor.numpy())
    return preds, [F.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]

# 使用matplotlib的subplots将一组images进行加载显示
def plot_classes_preds(net, images, labels):
    '''
    Generates matplotlib Figure using a trained network, along with images
    and labels from a batch, that shows the network's top prediction along
    with its probability, alongside the actual label, coloring this
    information based on whether the prediction was correct or not.
    Uses the "images_to_probs" function.
    '''
    preds, probs = images_to_probs(net, images)
    # plot the images in the batch, along with predicted and true labels
    fig = plt.figure(figsize=(12, 4))
    for idx in np.arange(4):
        ax = fig.add_subplot(1, 4, idx+1, xticks=[], yticks=[])
        matplotlib_imshow(images[idx], one_channel=True)
        ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
            classes[preds[idx]],
            probs[idx] * 100.0,
            classes[labels[idx]]),
                    color=("green" if preds[idx]==labels[idx].item() else "red"))
    return fig

# 计算准确率
def accuracy(net, images, labels):
    preds, probs = images_to_probs(net, images)
#     print(preds, labels)
#     print(preds == labels.numpy(), (preds == labels.numpy()).sum())
    
    acc = (preds == labels.numpy()).sum() / len(labels)
    return acc

# 计算在数据集上的loss、accuracy以及预测概率
def evaluate(net, dataloader):
    eval_loss = 0.0
    eval_acc = 0.0
    eval_probs = []
    eval_preds = []
    
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        outputs = net(inputs)
        loss = criterion(outputs, labels).item()
        acc = accuracy(net, inputs, labels).item()
        eval_loss += loss
        eval_acc += acc
        preds, probs = images_to_probs(net, inputs)
        eval_probs.append(probs)
        eval_preds.append(preds)
#         print(i, loss, acc)
        
    return eval_loss / len(dataloader), eval_acc / len(dataloader), np.concatenate(eval_probs, axis=0), np.concatenate(eval_preds, axis=0)

# 绘制各个类别的ROC曲线
def plot_roc_curve(probs, labels, classes, figsize=(40, 4)):
    fig, axes = plt.subplots(1, len(classes), figsize=figsize)
    
    labels = labels.numpy()
    for cls_index, cls_name in enumerate(classes, 0):
        # 如果分类结果中只有一类是不能绘制ROC曲线和计算AUC的
        cls_labels = labels.copy()
        # 将训练标签中当前类型的label置为1（正类），其余类型的label置为0（负类）
        cls_labels[labels != cls_index] = 0
        cls_labels[labels == cls_index] = 1
#         print(cls_labels[:50])
        
        # 在预测概率中，保留当前类别的预测概率（正类），负类的预测概率为1 - 当前类别的概率
        cls_probs = probs.copy()
        cls_probs[labels != cls_index] = 1 - cls_probs[labels != cls_index]
        cls_probs = cls_probs.copy()
#         print(cls_probs[:50])
        
        fpr, tpr, thr = roc_curve(cls_labels, cls_probs)
#         print(fpr[:50])
#         print(tpr[:50])
        axes[cls_index].plot(fpr, tpr, label="AUC: {:.2f}".format(auc(fpr, tpr)))
        axes[cls_index].set_xlabel("FPR")
        axes[cls_index].set_ylabel("TPR")
        axes[cls_index].set_title(classes[cls_index])
        axes[cls_index].legend(loc="lower right")
#     plt.show()
    return fig

# 绘制各个类别的PR曲线
def plot_pr_curve(writer, probs, labels, classes, step):
    labels = labels.numpy()
    for cls_index, cls_name in enumerate(classes, 0):
        # 如果分类结果中只有一类是不能绘制ROC曲线和计算AUC的
        cls_labels = labels.copy()
        # 将训练标签中当前类型的label置为1（正类），其余类型的label置为0（负类）
        cls_labels[labels != cls_index] = 0
        cls_labels[labels == cls_index] = 1
#         print(cls_labels[:50])
        
        # 在预测概率中，保留当前类别的预测概率（正类），负类的预测概率为1 - 当前类别的概率
        cls_probs = probs.copy()
        cls_probs[labels != cls_index] = 1 - cls_probs[labels != cls_index]
        cls_probs = cls_probs.copy()
#         print(cls_probs[:50])
        writer.add_pr_curve("PR/{}".format(cls_name), cls_labels, cls_probs, global_step=step)
    return writer

# 计算各个类别的准确率
def plot_accuracy(writer, prebs, labels, classes):
    labels = labels.numpy()
    acc_dict = {}
    for cls_index, cls_name in enumerate(classes, 0):
        cls_labels = labels[labels == cls_index]
        cls_prebs = prebs[labels == cls_index]

        acc = (cls_labels == cls_prebs).sum() / cls_labels.shape[0]
        acc_dict[cls_name] = acc
        
    for cls_name, cls_acc in acc_dict.items():
        writer.add_scalar("Accuracy/{}".format(cls_name), cls_acc, epoch * len(trainloader) + i)
    return writer

定义模型训练过程：

# 记录训练误差和训练准确率
running_loss = 0.0
running_acc = 0.0
for epoch in range(3):  # loop over the dataset multiple times

    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        running_acc += accuracy(net, inputs, labels).item()
        if i % 1000 == 999:    # every 1000 mini-batches...
            net.eval()
            # 记录训练集的整体损失和准确率
            # ...log the running loss
            writer.add_scalar('Loss/training loss', running_loss / 1000, epoch * len(trainloader) + i)
            writer.add_scalar("Accuracy/training accuracy", running_acc / 1000, epoch * len(trainloader) + i)
            
            # 记录测试集的整体损失和准确率
            eval_loss, eval_acc, eval_probs, eval_prebs = evaluate(net, testloader)
            writer.add_scalar("Loss/testing loss", eval_loss, epoch * len(trainloader) + i)
            writer.add_scalar("Accuracy/testing accuracy", eval_acc, epoch * len(trainloader) + i)
            
            # 记录测试集各个类别的准确率
#             for cls_name, cls_acc in plot_accuracy(eval_prebs, testset.targets, classes).items():
#                 writer.add_scalar("Accuracy/{}".format(cls_name), cls_acc, epoch * len(trainloader) + i)
            writer = plot_accuracy(writer, eval_prebs, testset.targets, classes)
            
            # 记录测试集各个类别的P-R曲线
            writer = plot_pr_curve(writer, eval_probs, testset.targets, classes, 
                                   step=epoch * len(trainloader) + i)
            
            # 记录测试集各个类别的ROC曲线
            writer.add_figure("testing roc curve", 
                              plot_roc_curve(eval_probs, testset.targets, classes), 
                              global_step=epoch * len(trainloader) + i)

            # 输出样本的真实标签和预测标签、预测概率
            # ...log a Matplotlib Figure showing the model's predictions on a
            # random mini-batch
            writer.add_figure('predictions vs. actuals',
                            plot_classes_preds(net, inputs, labels),
                            global_step=epoch * len(trainloader) + i)
            
            print("epoch: {}, iter: {}, train loss: {}, train acc: {}, eval loss: {}, eval acc: {}".format(epoch, i, 
                                                                                                           running_loss / 1000, 
                                                                                                           running_acc / 1000, 
                                                                                                           eval_loss, 
                                                                                                           eval_acc))
            running_loss = 0.0
            running_acc = 0.0
            net.train()
writer.close()
print('Finished Training')

tensorboard页面可以看到训练的损失函数、准确率变化曲线，模型预测的PR曲线和ROC曲线如下：

1）、损失函数曲线

2）、准确率曲线

3）、各个类别的准确率曲线

4）、各个类别的PR曲线

5）、各个类别的ROC曲线

参考链接

VISUALIZING MODELS, DATA, AND TRAINING WITH TENSORBOARD

TORCH.UTILS.TENSORBOARD

使用scikit-learn计算分类器的ROC曲线及AUC值