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。