30分钟吃掉wandb模型训练可视化

wandb是"我爱你，大baby"首字母的缩写。

顾名思义，她是炼丹师的大宝贝，是炼丹师最爱的炼丹伴侣。

公众号算法美食屋后台回复关键词：wandb，获取本教程 notebook源码 和 B站视频演示。

just kidding, 开个玩笑！

wandb全称weights&bias，是一款类似TensorBoard

的机器学习可视化分析工具。

相比TensorBoard，wandb具有如下主要优势：

• 日志上传云端永久存储，便于分享不怕丢失。

• 可以存管代码,数据集和模型的版本，随时复现。(wandb.Artifact)

• 可以使用交互式表格进行case分析(wandb.Table)

• 可以自动化模型调参。(wandb.sweep)

官方文档：https://docs.wandb.ai/

总体来说，wandb目前的核心功能有以下4个：

1，实验跟踪：experiment tracking （wandb.log）

2，版本管理：version management (wandb.log_artifact, wandb.save)

3，case分析：case visualization (wandb.Table, wandb.Image)

4，超参调优：model optimization (wandb.sweep)

本文我们主要介绍 前3个能力，超参调优的介绍在下一篇文章。

〇，注册wandb

使用wandb可视化模型训练过程需要在  https://wandb.ai/ 注册账户，

并在个人settings页面获取 API keys。

#import os
#os.environ["WANDB_API_KEY"] = "xxxx"

import wandb
wandb.login()

一，实验跟踪

wandb 提供了类似 TensorBoard的实验跟踪能力，主要包括：

• 模型配置超参数的记录

• 模型训练过程中loss，metric等各种指标的记录和可视化

• 图像的可视化(wandb.Image)

• 其他各种Media(wandb.Vedio, wandb.Audio, wandb.Html, 3D点云等)

import os,PIL 
import numpy as np
from torch.utils.data import DataLoader, Dataset
import torch 
from torch import nn 
import torchvision 
from torchvision import transforms
import datetime
import wandb 
from argparse import Namespace

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
config = Namespace(
    project_name = 'wandb_demo',
    
    batch_size = 512,
    
    hidden_layer_width = 64,
    dropout_p = 0.1,
    
    lr = 1e-4,
    optim_type = 'Adam',
    
    epochs = 15,
    ckpt_path = 'checkpoint.pt'
)

def create_dataloaders(config):
    transform = transforms.Compose([transforms.ToTensor()])
    ds_train = torchvision.datasets.MNIST(root="./mnist/",train=True,download=True,transform=transform)
    ds_val = torchvision.datasets.MNIST(root="./mnist/",train=False,download=True,transform=transform)

    ds_train_sub = torch.utils.data.Subset(ds_train, indices=range(0, len(ds_train), 5))
    dl_train =  torch.utils.data.DataLoader(ds_train_sub, batch_size=config.batch_size, shuffle=True,
                                            num_workers=2,drop_last=True)
    dl_val =  torch.utils.data.DataLoader(ds_val, batch_size=config.batch_size, shuffle=False, 
                                          num_workers=2,drop_last=True)
    return dl_train,dl_val

def create_net(config):
    net = nn.Sequential()
    net.add_module("conv1",nn.Conv2d(in_channels=1,out_channels=config.hidden_layer_width,kernel_size = 3))
    net.add_module("pool1",nn.MaxPool2d(kernel_size = 2,stride = 2)) 
    net.add_module("conv2",nn.Conv2d(in_channels=config.hidden_layer_width,
                                     out_channels=config.hidden_layer_width,kernel_size = 5))
    net.add_module("pool2",nn.MaxPool2d(kernel_size = 2,stride = 2))
    net.add_module("dropout",nn.Dropout2d(p = config.dropout_p))
    net.add_module("adaptive_pool",nn.AdaptiveMaxPool2d((1,1)))
    net.add_module("flatten",nn.Flatten())
    net.add_module("linear1",nn.Linear(config.hidden_layer_width,config.hidden_layer_width))
    net.add_module("relu",nn.ReLU())
    net.add_module("linear2",nn.Linear(config.hidden_layer_width,10))
    net.to(device)
    return net 

def train_epoch(model,dl_train,optimizer):
    model.train()
    for step, batch in enumerate(dl_train):
        features,labels = batch
        features,labels = features.to(device),labels.to(device)

        preds = model(features)
        loss = nn.CrossEntropyLoss()(preds,labels)
        loss.backward()

        optimizer.step()
        optimizer.zero_grad()
    return model

def eval_epoch(model,dl_val):
    model.eval()
    accurate = 0
    num_elems = 0
    for batch in dl_val:
        features,labels = batch
        features,labels = features.to(device),labels.to(device)
        with torch.no_grad():
            preds = model(features)
        predictions = preds.argmax(dim=-1)
        accurate_preds =  (predictions==labels)
        num_elems += accurate_preds.shape[0]
        accurate += accurate_preds.long().sum()

    val_acc = accurate.item() / num_elems
    return val_acc

def train(config = config):
    dl_train, dl_val = create_dataloaders(config)
    model = create_net(config); 
    optimizer = torch.optim.__dict__[config.optim_type](params=model.parameters(), lr=config.lr)
    #======================================================================
    nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    wandb.init(project=config.project_name, config = config.__dict__, name = nowtime, save_code=True)
    model.run_id = wandb.run.id
    #======================================================================
    model.best_metric = -1.0
    for epoch in range(1,config.epochs+1):
        model = train_epoch(model,dl_train,optimizer)
        val_acc = eval_epoch(model,dl_val)
        if val_acc>model.best_metric:
            model.best_metric = val_acc
            torch.save(model.state_dict(),config.ckpt_path)   
        nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"epoch【{epoch}】@{nowtime} --> val_acc= {100 * val_acc:.2f}%")
        #======================================================================
        wandb.log({'epoch':epoch, 'val_acc': val_acc, 'best_val_acc':model.best_metric})
        #======================================================================        
    #======================================================================
    wandb.finish()
    #======================================================================
    return model

model = train(config) ##3,2,1 点火🔥🔥

二，版本管理

除了可以记录实验日志传递到 wandb 网站的云端服务器 并进行可视化分析。

wandb还能够将实验关联的数据集，代码和模型 保存到 wandb 服务器。

非常便于我们以后 或者其他人 对实验结果进行复现。

我们可以通过 wandb.log_artifact的方法来保存任务的关联的重要成果。

例如dataset, code，和 model，并进行版本管理。

注：artifact翻译为"工件"，是指软件开发中产出的最终成果.

我们先使用run_id 恢复 run任务，以便继续记录。

#resume the run 
import wandb 

run = wandb.init(project='wandb_demo', id= model.run_id, resume='must')

# save dataset 
arti_dataset = wandb.Artifact('mnist', type='dataset')
arti_dataset.add_dir('mnist/')
wandb.log_artifact(arti_dataset)

# save code 

arti_code = wandb.Artifact('ipynb', type='code')
arti_code.add_file('./30分钟吃掉wandb可视化模型分析.ipynb')
wandb.log_artifact(arti_code)

# save model

arti_model = wandb.Artifact('cnn', type='model')
arti_model.add_file(config.ckpt_path)
wandb.log_artifact(arti_model)

wandb.finish() #finish时会提交保存

三，Case分析

利用 wandb.Table 我们可以在 wandb的 dashboard 进行交互式可视化的 case分析。

#resume the run 
import wandb 
run = wandb.init(project=config.project_name, id= model.run_id, resume='must')

import matplotlib.pyplot as plt 

transform = transforms.Compose([transforms.ToTensor()])
ds_train = torchvision.datasets.MNIST(root="./mnist/",train=True,download=True,transform=transform)
ds_val = torchvision.datasets.MNIST(root="./mnist/",train=False,download=True,transform=transform)
    
# visual the  prediction
device = None
for p in model.parameters():
    device = p.device
    break

plt.figure(figsize=(8,8)) 
for i in range(9):
    img,label = ds_val[i]
    tensor = img.to(device)
    y_pred = torch.argmax(model(tensor[None,...])) 
    img = img.permute(1,2,0)
    ax=plt.subplot(3,3,i+1)
    ax.imshow(img.numpy())
    ax.set_title("y_pred = %d"%y_pred)
    ax.set_xticks([])
    ax.set_yticks([]) 
plt.show()

def data2fig(data):
    import matplotlib.pyplot as plt 
    fig = plt.figure()
    ax = fig.add_subplot()
    ax.imshow(data)
    ax.set_xticks([])
    ax.set_yticks([]) 
    return fig

def fig2img(fig):
    import io,PIL
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    img = PIL.Image.open(buf)
    return img

from tqdm import tqdm 
good_cases = wandb.Table(columns = ['Image','GroundTruth','Prediction'])
bad_cases = wandb.Table(columns = ['Image','GroundTruth','Prediction'])

# 找到50个good cases 和 50 个bad cases

plt.close()

for i in tqdm(range(1000)):
    features,label = ds_val[i]
    tensor = features.to(device)
    y_pred = torch.argmax(model(tensor[None,...])) 
    
    # log badcase
    if y_pred!=label:
        if len(bad_cases.data)<50:
            data = features.permute(1,2,0).numpy()
            input_img = wandb.Image(fig2img(data2fig(data)))
            bad_cases.add_data(input_img,label,y_pred)
            
    # log goodcase
    else:
        if len(good_cases.data)<50:
            data = features.permute(1,2,0).numpy()
            input_img = wandb.Image(fig2img(data2fig(data)))
            good_cases.add_data(input_img,label,y_pred)

wandb.log({'good_cases':good_cases,'bad_cases':bad_cases})

wandb.finish()

然后就可以在 DashBoard中对Table 进行交互式的分析了，

包括常用的 sort, filter, group 等操作，妥妥的一个网页版的 Excel。