Bootstrap

机器学习实用工具 Wandb(1)—— 实验追踪

  • 在做机器学习项目时,比如这个典型例子,常常遇到以下几个痛点
    1. 记录训练曲线的代码繁琐,与模型代码耦合度高,观感差又不好修改
    2. 自己做可视化效果较差,要做好又太浪费时间
    3. 调参时各种超参数模型难以管理,不易进行性能比较,网格搜索代码也很麻烦
  • wandb 是一个实验记录平台,它可以快速实现美观的可视化效果,支持多种机器学习框架,代码侵入小,还能帮助我们进行超参参数搜索及版本控制,可以有效解决以上问题。如果使用团队版本,还可进行团队内实验结果共享,实验环境同步,实验报告制作等功能。它主要有以下四个组件构成:
    1. Dashboard: 实验跟踪
    2. Artifacts: 数据集版本控制、模型版本控制
    3. Sweeps: 超参数优化
    4. Reports: 保存和共享可重现的结果
  • 本文主要介绍如何集成 wandb + pytorch 来记录实验过程并可视化,参考文档(注:英文文档比中文文档更细致)
    1. Quickstart
    2. Experiment Tracking
    3. Simple_PyTorch_Integration

1. 安装和注册

  1. 直接 pip install wandb 安装
  2. wandb 官网注册账号,注意现在只有注册成个人使用才免费
    在这里插入图片描述
    注册好之后复制给出的私钥,然后在命令行执行 wandb login,根据提示输入私钥,即可建立起本地环境和 wandb 平台的联系

2. 跟踪训练过程

  • 简单概括下工作流程:将 wandb 集成到需要观测的机器学习代码后,它会在本地建立一个仓库记录所有指定的实验数据,同时数据被异步上传到 wandb 服务器,并在网页端实时可视化。如果没有网络或者数据涉密,也可以建立本地服务器来完成全部流程

  • Wandb 组成结构如下
    在这里插入图片描述

    1. entity 就是 wandb 账户,可以是个人或组织,它管理一个用户生成的所有 log 记录
    2. wandb 日志的最大组成单元是 project,涵盖一个项目要记录的所有训练曲线、数据集和模型版本、代码副本等等;
    3. wandb 日志的最小组成单元是 run,它是一个 wandb project 的基本组成单位。一个 run 本质上是对一段训练或评估过程的记录,体现为一组数据曲线以及其他所有使用 wandb.log 方法记录的信息(图像、视频、文本等),每次设定新的随机种子或超参数的训练都会生成一个 run。另外,run 也是 wandb 网页端可视化的最小单位,相关说明参考这里
    4. 可以用 group 组织多个 run,这些被组织的 run 的 loss & metric 曲线将在网页端显示为一条 “平均线”,即前景为 mean 曲线,背景为标准差阴影。一个 run 所属的 group 可以在调用 wandb.init 时通过 group 参数设置,也可以在网页端手动选择 wandb.config 中的任意字段作为 group 标签,相关说明参考这里

2.1 在代码中集成 wandb

  • 将 wandb 集成到典型的 ML pipeline 中的伪代码如下
    # import the library
    import wandb
    
    # 1. start a new experiment
    wandb.init(project="new-sota-model")
    
    # 2. capture a dictionary of hyperparameters with config
    wandb.config = {"learning_rate": 0.001, "epochs": 100, "batch_size": 128}
    
    # set up model and data
    model, dataloader = get_model(), get_data()
    
    # 3. optional: track gradients
    wandb.watch(model)
    
    for batch in dataloader:
    	 metrics = model.training_step()
    	 # 4. log metrics inside your training loop to visualize model performance
    	 wandb.log(metrics)
    
    # 5. Log an artifact to W&B
    wandb.log_artifact(model)
    
    # 6. optional: save model at the end
    model.to_onnx()
    wandb.save("model.onnx")
    
  • 这里的关键代码只有 6 条,对模型代码入侵很小,下面依次介绍

2.1.1 wandb.init()

  • wandb.init():在训练或评估过程开始之前调用,它会新建一个 run,并且创建一个本地目录用于保存所有日志和文件。后续记录 metric 信息或保存模型文件时,这些数据会存储到上述本地目录,同时异步上传到 wandb 服务器,在网页端实时呈现出来
  • 该方法原型如下
    def init(
        job_type: Optional[str] = None,					# 该 run 的 job 类型,一个 group 内可以有多个 job(如 train 和 eval),用来对 run 进行过滤和分组
        dir: Optional[StrPath] = None,					# 保存该 run 信息的本地目录
        config: Union[Dict, str, None] = None,			# 超参数 & 元数据字典,可以在此设置下一节的 wandb.config
        project: Optional[str] = None,					# 该 run 所属的 project 名称
        entity: Optional[str] = None,					# 该 run 所属的组织名,默认为用户账号
        reinit: Optional[bool] = None,
        tags: Optional[Sequence] = None,				# 字符串列表,可以用来对 run 进行过滤和分组
        group: Optional[str] = None,					# 该 run 所属的 group 名称,group 可以在一个 project 内组织多个 run
        name: Optional[str] = None,						# 该 run 的名称
        notes: Optional[str] = None,					# 对该 run 的较长描述,它会和 config 一起以表格形式在网页端显示
        magic: Optional[Union[dict, str, bool]] = None,
        config_exclude_keys: Optional[List[str]] = None,
        config_include_keys: Optional[List[str]] = None,
        anonymous: Optional[str] = None,
        mode: Optional[str] = None,
        allow_val_change: Optional[bool] = None,
        resume: Optional[Union[bool, str]] = None,
        force: Optional[bool] = None,
        tensorboard: Optional[bool] = None,
        sync_tensorboard: Optional[bool] = None,		# 从tensorboard 同步 log 信息并保存相关事件文件
        monitor_gym: Optional[bool] = None,				# 使用OpenAI Gym时自动记录环境视频
        save_code: Optional[bool] = None,				# 是否将 main 脚本保存到服务器,对于复现结果有益
        id: Optional[str] = None,						# 该 run 的唯一标识
        settings: Union[Settings, Dict[str, Any], None] = None,
    ) -> Union[Run, RunDisabled, None]:
    
    其中我注释的是相对常用的,完整参数说明参考官方文档
  • 一个 run 自 wandb.init() 开始,至 wandb.finish() 或程序退出结束(退出时自动调用 wandb.finish),另外也可以用 python 的 with 语法明确一个 run 的起止范围
    import wandb
    
    # 手动开启和关闭一个 run
    wandb.init()
    wandb.finish()
    assert wandb.run is None
    
    # 用with语法自动设置 run 的范围
    with wandb.init() as run:
        pass  # log data here
    assert wandb.run is None
    
    如果我们想在一次运行中创建多个 run(比如测试多个随机种子),则需要多次调用以上结构

2.1.2 wandb.config

  • wandb.config:这个字典对象用来保存实验设置,包括
    1. hyperparameters 超参数,这些参数会影响模型性能,后续需要进行网格搜索调整;
    2. metadata 元数据,包括数据集名称、模型类型等实验信息,它们对分析实验和复现实验很有用
  • 一些值得注意的点
    1. 在 wandb 网页端可以通过各种 config 参数值对所有 Run 进行分组,方便我们比较不同的设置如何影响模型性能
    2. 注意 loss 和 metric 等变化量不应写入 wandb.config 中
    3. wandb.config 通常作为 wandb.init 的参数进行设置,也可以像本章开头的示例那样单独设置,设置之后可以通过 wandb.config['key'] 形式访问参数值或更新其值
    4. wandb.config 中允许设置嵌套字典,它们会被自动转换为 A.a 形式
  • 可以直接用 argparse 设置 wandb.config,这种做法非常常见,示例如下
    # config_experiment.py
    import wandb
    import argparse
    import numpy as np 
    import random
    
    
    # Training and evaluation demo code
    def train_one_epoch(epoch, lr, bs): 
        acc = 0.25 + ((epoch/30) +  (random.random()/10))
        loss = 0.2 + (1 - ((epoch-1)/10 +  random.random()/5))
        return acc, loss
    
    def evaluate_one_epoch(epoch): 
        acc = 0.1 + ((epoch/20) +  (random.random()/10))
        loss = 0.25 + (1 - ((epoch-1)/10 +  random.random()/6))
        return acc, loss
    
    
    def main(args):
        # Start a W&B Run
        run = wandb.init(project="config_example", config=args)
    
        # Access values from config dictionary and store them 
        # into variables for readability
        lr  =  wandb.config['learning_rate']
        bs = wandb.config['batch_size']
        epochs = wandb.config['epochs']
    
        # Simulate training and logging values to W&B 
        for epoch in np.arange(1, epochs):
            train_acc, train_loss = train_one_epoch(epoch, lr, bs)
            val_acc, val_loss = evaluate_one_epoch(epoch)
    
            wandb.log({
            'epoch': epoch, 
            'train_acc': train_acc,
            'train_loss': train_loss, 
            'val_acc': val_acc, 
            'val_loss': val_loss
            })
    
    if __name__ == "__main__":
      parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
      parser.add_argument(
        "-b",
        "--batch_size",
        type=int,
        default=32,
        help="Batch size")
      parser.add_argument(
        "-e",
        "--epochs",
        type=int,
        default=50,
        help="Number of training epochs")
      parser.add_argument(
        "-lr",
        "--learning_rate",
        type=int,
        default=0.001,
        help="Learning rate")    
    
    
    args = parser.parse_args()
    main(args)
    
  • 另外也可以用 yaml 文件设置 wandb.config,只要创建一个名为 config-defaults.yaml 的文件,键值对就会自动传递给 wandb.config,一个 yaml 配置文件示例如下
    # config-defaults.yaml
    # sample config defaults file
    epochs:
      desc: Number of epochs to train over
      value: 100
    batch_size:
      desc: Size of each mini-batch
      value: 32
    
    可以使用命令行参数 --configs 加载不同的配置文件;在 wandb.init 中传入 config 参数可以覆盖这种默认值;另外还可以像下面这样混合字典形式参数和yaml文件形式参数
    hyperparameter_defaults = dict(
        dropout=0.5,
        batch_size=100,
        learning_rate=0.001,
        )
    
    config_dictionary = dict(
        yaml=my_yaml_file,
        params=hyperparameter_defaults,
        )
    
    wandb.init(config=config_dictionary)
    

2.1.3 wandb.watch()

  • wandb.watch():在训练开始前调用,调用后会按固定的 batch 周期记录模型的梯度和参数

    如果开启了记录,每次 Run 的梯度和参数都会记录在 wandb 网页端,如
    在这里插入图片描述 在这里插入图片描述

    这些数据非常有用,可以帮助我们判断是否发生了梯度爆炸/梯度消失等问题,也能帮助我们判断是否有极端参数值主导模型输出,从而决定需要增加正则化项或 dropout

2.1.4 wandb.log()

  • wandb.log():在循环中周期性调用来记录各项数据指标,每次调用时会向 history 对象追加一个新记录,并更新 summary 对象。 history 对象是一组像字典一样的对象,记录了各项指标随时间的变化,可以显示为折线图;summary 对象默认记录的是最后一次 wandb.log() 记的值,也可手动设定为记录 history 的某种统计信息,比如最高精度或最低损失等,wandb 网站会自动利用这些信息进行 Run 的排序
  • 另外,log 方法也可上传图像、视频、html 等多种格式,可以参考 wandb使用教程(一):基础用法 以及官方文档

2.1.5 wandb.log_artifact()

  • Artifacts 可以将任何序列化数据作为 run 的输入和输出进行跟踪和版本控制。例如
    1. 进行模型训练的 run 将数据集作为输入,将训练好的模型作为输出
    2. 可以执行一些特殊的 run,将数据集作为输入,将模型 checkpoints 作为输出
  • Artifacts 总是能告诉我们 “这个模型是在我的哪个版本的数据集上训练的” ,它的使用场景主要包括
    1. 查看模型的来源,以及训练它使用的数据
    2. 记录每个数据集的更改或模型 checkpoint 的版本
    3. 方便地在团队中重用模型和数据集
  • 下面给出一个创建 Artifacts 的最小示例
    # 创建一个run
    run = wandb.init(project="artifacts-example", job_type="add-dataset")
    
    # 创建artifact对象
    artifact = wandb.Artifact(name="my_data", type="dataset")
    
    # 向artifact对象添加一个或多个文件,比如模型文件或数据集
    artifact.add_dir(local_path="./dataset.h5")  	# Add dataset directory to artifact
    
    # 将artifact记录到wandb。
    run.log_artifact(artifact)  					# Logs the artifact version "my_data:v0"
    

2.1.6 wandb.save()

  • wandb.save():在实验完成后调用这个来生成并保存 onnx 格式的模型,这是一种表示机器学习模型的通用开源格式,常常用来将 pytorch 模型转换为 TF 或 Keras 模型
  • 保存的 .onnx 模型会被自动同步到 wandb 网站上,网站内嵌了 onnx 可视化工具,可以给出漂亮的网络结构图,例如下面这个两层 MLP
    在这里插入图片描述

2.2 FashionMNIST 分类示例

  • 下面我们将 wandb 嵌入到之前 经典机器学习方法(3)—— 多层感知机 中介绍过的使用两层 MLP 做 FishonMIST 分类任务的代码上,来观察不同隐藏层尺寸对性能的影响
  • 首先把上文中最后的 pytorch 代码整理为符合经典 ML pipeline 的结构,如下
    import torch
    import torchvision
    import torchvision.transforms as transforms
    from torch import nn
    from torch.nn import init
    import random
    import numpy as np
    from tqdm import tqdm
    import argparse
    from pathlib import Path
    import os
    
    def model_pipeline(hyperparameters):
        '''
        the overall pipeline, which is pretty typical for model-training
        '''
        config=hyperparameters
        for seed in hyperparameters.seeds:
            set_random_seed(seed)
            
            # make the model, data, and optimization problem
            model, train_loader, val_loader, test_loader, loss, optimizer = make(config)
            print(model)
    
            # and use them to train the model
            train(model, train_loader, val_loader, loss, optimizer, config)
    
            # and test its final performance
            test(model, test_loader)
    
    def set_random_seed(random_seed):
        torch.backends.cudnn.deterministic = True
        random.seed(random_seed)         
        np.random.seed(random_seed)
        torch.manual_seed(random_seed)
        torch.cuda.manual_seed_all(random_seed)
    
    def make(config):
        '''
        make the data, model, loss and optimizer
        '''
        # Make the data
        train = torchvision.datasets.FashionMNIST(root='./Datasets/FashionMNIST', train=True, transform=transforms.ToTensor(), download=True)
        test = torchvision.datasets.FashionMNIST(root='./Datasets/FashionMNIST', train=False, transform=transforms.ToTensor(), download=True)
    
        train_dataset = torch.utils.data.Subset(train, indices=range(0, int(0.8*len(train))))
        val_dataset = torch.utils.data.Subset(train, indices=range(int(0.8*len(train)), len(train)))
        test_dataset = torch.utils.data.Subset(test, indices=range(0, len(test), 1))
    
        train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=config.batch_size, shuffle=True, pin_memory=True, num_workers=4)
        val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=len(val_dataset), shuffle=True, pin_memory=True, num_workers=4) 
        test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=config.batch_size, shuffle=True, pin_memory=True, num_workers=4)
    
        # Make the model
        model = MLP(784, 10, config.num_hiddens).to(device)
        for params in model.parameters():
            init.normal_(params, mean=0, std=0.01)
            
        # Make the loss and optimizer
        loss = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)
        
        return model, train_loader, val_loader, test_loader, loss, optimizer
    
    class FlattenLayer(nn.Module):
        '''
        这个自定义 Module 将二维的图像输入拉平成一维向量
        '''
        def __init__(self):
            super(FlattenLayer, self).__init__()
            
        def forward(self, x): # x shape: (batch, *, *, ...)
            return x.view(x.shape[0], -1)
    
    def MLP(num_inputs, num_outputs, num_hiddens):
        model = nn.Sequential(
            FlattenLayer(),
            nn.Linear(num_inputs, num_hiddens),
            nn.ReLU(),
            nn.Linear(num_hiddens, num_outputs), 
        )
        return model
    
    def train(model, train_loader, val_loader, loss, optimizer, config):
    
        # Run training 
        total_batches = len(train_loader) * config.epochs
        example_cnt = 0  # number of examples seen
        batch_cnt = 0
        for epoch in range(config.epochs):
            with tqdm(total=len(train_loader), desc=f'epoch {epoch+1}') as pbar:   
                for _, (images, labels) in enumerate(train_loader):
    
                    train_loss = train_batch(images, labels, model, optimizer, loss)
                    example_cnt +=  len(images)
                    batch_cnt += 1
    
                    # Report metrics every 20th batch
                    if (batch_cnt + 1) % 20 == 0:
                        val_accuracy, val_loss = validation(model, val_loader, loss)
    
                        # update tqdm information
                        pbar.set_postfix({
                            'val_acc':
                            '%.3f' % val_accuracy,
                            'val_loss':
                            '%.3f' % val_loss,
                        })
    
                    pbar.update(1)
                        
    def train_batch(images, labels, model, optimizer, loss):
        images, labels = images.to(device), labels.to(device)
        
        # Forward pass
        outputs = model(images)
        train_loss = loss(outputs, labels)
        
        # Backward pass
        optimizer.zero_grad()
        train_loss.backward()
    
        # Step with optimizer
        optimizer.step()
        return train_loss
    
    def test(model, test_loader):
        model.eval()
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                test_accuracy = correct/total
    
            print(f"Accuracy of the model on the {total} test images: {test_accuracy:%}")
        
        model.train()
        return test_accuracy
    
    def validation(model, val_loader, loss):
        model.eval()
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                val_accuracy = correct/total
                val_loss = loss(outputs, labels)
    
        model.train()
        return val_accuracy, val_loss
    
    if __name__ == '__main__':
        # random seeds
        random_seeds = (43,44,45)
        
        # Device configuration
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(device)
    
        # as the start of our workflow, store hyperparameters
        parser = argparse.ArgumentParser()
        parser.add_argument('--seeds', type=int, default=random_seeds)
        parser.add_argument('--epochs', type=int, default=5)  
        parser.add_argument('--batch_size', type=int, default=512) 
        parser.add_argument('--learning_rate', type=float, default=0.1)
        parser.add_argument('--num_hiddens', type=int, default=32)
        
        args = parser.parse_args()
        config = args
    
        # Build, train and analyze the model with the pipeline
        model = model_pipeline(config)
    
  • 接下来按 2.1 节说明增加 wandb 方法,完整代码如下,请读者自行对比
    import torch
    import torchvision
    import torchvision.transforms as transforms
    from torch import nn
    from torch.nn import init
    import random
    import numpy as np
    import wandb 
    from tqdm import tqdm
    import argparse
    from pathlib import Path
    import os
    
    def model_pipeline(hyperparameters):
        '''
        the overall pipeline, which is pretty typical for model-training
        '''
        # set the location where all the data logged from script will be saved, which will be synced to the W&B cloud
        # the default location is ./wandb
        run_dir = Path(f"{os.getcwd()}/wandb_local") / hyperparameters.project_name / hyperparameters.experiment_name
        if not run_dir.exists():
            os.makedirs(str(run_dir))
    
        for seed in hyperparameters.seeds:
            set_random_seed(seed)
    
            # tell wandb to get started
            with wandb.init(config=vars(hyperparameters),
                            project=hyperparameters.project_name,
                            group=hyperparameters.scenario_name,
                            name=hyperparameters.experiment_name+"_"+str(seed),
                            notes=hyperparameters.note, 
                            dir=run_dir):
    
                # access all HPs through wandb.config, ensuring the values you chose and logged are always the ones that get used in your model
                config = wandb.config
    
                # make the model, data, and optimization problem
                model, train_loader, val_loader, test_loader, loss, optimizer = make(config)
                print(model)
    
                # and use them to train the model
                train(model, train_loader, val_loader, loss, optimizer, config)
    
                # and test its final performance
                test(model, test_loader)
    
                # Save the model in the exchangeable ONNX format
                # Passing that filename to wandb.save ensures that the model parameters are saved to W&B's servers: 
                torch.onnx.export(model, torch.randn(config.batch_size, 1, 28, 28).to(device), "model.onnx")
                wandb.save("model.onnx")
                wandb.finish()
    
    def set_random_seed(random_seed):
        torch.backends.cudnn.deterministic = True
        random.seed(random_seed)         
        np.random.seed(random_seed)
        torch.manual_seed(random_seed)
        torch.cuda.manual_seed_all(random_seed)
    
    def make(config):
        '''
        make the data, model, loss and optimizer
        '''
        # Make the data
        train = torchvision.datasets.FashionMNIST(root='./Datasets', train=True, transform=transforms.ToTensor(), download=True)
        test = torchvision.datasets.FashionMNIST(root='./Datasets', train=False, transform=transforms.ToTensor(), download=True)
    
        train_dataset = torch.utils.data.Subset(train, indices=range(0, int(0.8*len(train))))
        val_dataset = torch.utils.data.Subset(train, indices=range(int(0.8*len(train)), len(train)))
        test_dataset = torch.utils.data.Subset(test, indices=range(0, len(test), 1))
    
        train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=config.batch_size, shuffle=True, pin_memory=True, num_workers=4)
        val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=len(val_dataset), shuffle=True, pin_memory=True, num_workers=4) 
        test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=config.batch_size, shuffle=True, pin_memory=True, num_workers=4)
    
        # Make the model
        model = MLP(784, 10, config.num_hiddens).to(device)
        for params in model.parameters():
            init.normal_(params, mean=0, std=0.01)
            
        # Make the loss and optimizer
        loss = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)
        
        return model, train_loader, val_loader, test_loader, loss, optimizer
    
    class FlattenLayer(nn.Module):
        '''
        这个自定义 Module 将二维的图像输入拉平成一维向量
        '''
        def __init__(self):
            super(FlattenLayer, self).__init__()
            
        def forward(self, x): # x shape: (batch, *, *, ...)
            return x.view(x.shape[0], -1)
    
    def MLP(num_inputs, num_outputs, num_hiddens):
        model = nn.Sequential(
            FlattenLayer(),
            nn.Linear(num_inputs, num_hiddens),
            nn.ReLU(),
            nn.Linear(num_hiddens, num_outputs), 
        )
        return model
    
    def train(model, train_loader, val_loader, loss, optimizer, config):
        # wandb.watch will log the gradients and the parameters of your model, every log_freq steps of training.
        # it need to be called before start training
        wandb.watch(model, loss, log="all", log_freq=10)
    
        # Run training and track with wandb
        total_batches = len(train_loader) * config.epochs
        example_cnt = 0  # number of examples seen
        batch_cnt = 0
        for epoch in range(config.epochs):
            with tqdm(total=len(train_loader), desc=f'epoch {epoch+1}') as pbar:   # tqdm的进度条功能
                for _, (images, labels) in enumerate(train_loader):
    
                    train_loss = train_batch(images, labels, model, optimizer, loss)
                    example_cnt +=  len(images)
                    batch_cnt += 1
    
                    # Report metrics every 200th batch
                    if (batch_cnt + 1) % 20 == 0:
                        val_accuracy, val_loss = validation(model, val_loader, loss)
    
                        # update tqdm information
                        pbar.set_postfix({
                            'val_acc':
                            '%.3f' % val_accuracy,
                            'val_loss':
                            '%.3f' % val_loss,
                        })
    
                        # log the metrics to wandb
                        wandb.log({"epoch": epoch + 1, 
                                    "train_loss": train_loss, 
                                    'val_accuracy': val_accuracy, 
                                    'val_loss': val_loss}, 
                                    step=example_cnt)
    
                    pbar.update(1)
                        
    def train_batch(images, labels, model, optimizer, loss):
        images, labels = images.to(device), labels.to(device)
        
        # Forward pass 
        outputs = model(images)
        train_loss = loss(outputs, labels)
        
        # Backward pass ⬅
        optimizer.zero_grad()
        train_loss.backward()
    
        # Step with optimizer
        optimizer.step()
        return train_loss
    
    def test(model, test_loader):
        model.eval()
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                test_accuracy = correct/total
    
            wandb.log({'test_accuracy': test_accuracy})
            print(f"Accuracy of the model on the {total} test images: {test_accuracy:%}")
        
        model.train()
        return test_accuracy
    
    def validation(model, val_loader, loss):
        model.eval()
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                val_accuracy = correct/total
                val_loss = loss(outputs, labels)
    
        model.train()
        return val_accuracy, val_loss
    
    if __name__ == '__main__':
        # random seeds
        random_seeds = (43,44,45)
        
        # Device configuration
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(device)
    
        # as the start of our workflow, store hyperparameters and metadata in a config dictionary
        parser = argparse.ArgumentParser()
        parser.add_argument('--seeds', type=int, default=random_seeds)
        parser.add_argument('--epochs', type=int, default=5)  
        parser.add_argument('--batch_size', type=int, default=512) 
        parser.add_argument('--learning_rate', type=float, default=0.1)
        parser.add_argument('--num_hiddens', type=int, default=32)
    
        parser.add_argument('--dataset', type=str, default='FashionMNIST')
        parser.add_argument('--architecture', type=str, default='MLP')  
        parser.add_argument('--note', type=str, default='add some note for the run here')
        parser.add_argument('--project_name', type=str, default='Wandb_ExpTracking')
        parser.add_argument('--scenario_name', type=str, default='MLP_Hiddens')
        parser.add_argument('--experiment_name', type=str, default='seed')
        
        args = parser.parse_args()
        config = args
    
        # Build, train and analyze the model with the pipeline
        model = model_pipeline(config)
    

3. 可视化效果

  • 首次执行 2.2 节代码后,在你的 wandb 主页上就会出现名为 Wandb_ExpTracking 的 project。多次修改 num_hiddens 取值重复执行,然后在主页 project 栏找到 Wandb_ExpTracking 点进去,就会如下显示所有 Run 的平均性能
    在这里插入图片描述
    1. 每一张 Chart 记录的指标都是我们在代码中调用 wandb.log() 时传入的指标之一
    2. 每一张 Chart 右上角都可以点进去调整图像显示,比如调整曲线粗细和平滑度、设置 x 轴等
    3. 左边的项目结构是由我们设定的 metadata 决定,最后一级都是独立的 run,倒数第二级是这些 Run 的平均性能,会显示为一条曲线。这里我设置为最后一级是不同的随机种子,倒数第二级是隐藏层尺寸
    4. 下面一点 System 栏记录了训练过程中的硬件使用情况,这个也非常有用,可以帮助我们确定 batch_size 大小、确定速度瓶颈等
      在这里插入图片描述
  • 最左边的竖栏从上到下是
    1. Overview:显示项目的基础信息,在团队版本比较有用
    2. Workspace(当前位置)
    3. Table:显示所有 run 的统计数据
      在这里插入图片描述
      上面那个紫色的按钮可以点进去,根据不同指标的取值灵活设置项目结构,无论结构如何,最后一级都是独立的 Run,倒数第二级会自动变成这些 Run 的平均,显示在上面的曲线图中
    4. Reports:可以直接调用前面的各种插图写报告,不做介绍
    5. Sweeps:用来做超参数搜索的,以后介绍
    6. Artifacts:用来做模型版本控制的,以后介绍
  • 随便在 Workspace 找一个 Run 点进去,左边竖栏又会有五个
    1. Overview:显示这次 Run 的 metadata、指标,以及执行这个 run 的人员和软硬件环境等信息
    2. Charts:该 Run 的指标 Chart,和上面一致
    3. System:记录此 Run 执行过程的硬件资源占用情况
    4. Logs:记录此 Run 的 wandb.init 周期内所有的终端显示
    5. Files:记录此 Run 的虚拟环境、网络结构等信息
      在这里插入图片描述
      其中 model.onxx 可以点进去查看网络结构图,requirements.txt 记录了使用的所有依赖库,可以直接在本地重建虚拟环境
  • 关于客户端UI还可以参考官方文档

4. 与 wandb 联合使用

  • Tensorboard 比较轻量级,适合本地使用;wandb 更偏云端存储,适合团队间共享信息,二者也可以组合使用

  • 组合使用时,wandb 可以直接读取 Tensorboard 的本地仓库并同步所有数据,示例代码如下

    from torch.utils.tensorboard import SummaryWriter
    import wandb
    
    exp_name = f"tensorboard-wandb"
    
    # setup Wandb
    wandb.init(
        project='wandb_usage',
        entity=None,            # username
        sync_tensorboard=True,	# 同步 Tensorboard 数据
        config={'exp_name': exp_name},
        name=exp_name,
        save_code=True,         # 上传代码副本
    )
    
    # setup TensorBoard
    writer = SummaryWriter(f"runs/{exp_name}")
    for i in range(100):
        writer.add_scalar('test_loss', i*2, global_step=i)
    

    有可能出现无法同步 chart 数据,只能上传 GPU 等数据的情况,注意以下两个要点

    1. 先初始化 wandb,后初始化 SummaryWriter
    2. 以管理员身份运行(cmd 或 vscode),否则在 wandb/logs/debug-internal.log 中会看到以下错误
      Thread-30 :16652 [tb_watcher.py:_process_events():269] Encountered tensorboard directory watcher error: [WinError 1314] A required privilege is not held by the client: 
      
  • 更多联合使用的说明,参考 wandb 官方文档:Integrations/TensorBoard

;