【Python】科研代码学习：十四 wandb (可视化AI工具)

wandb 介绍

• 【wandb官网】
  wandb 是 Weights & Biases 的缩写（w and b）
• 核心作用：
  • 可视化重要参数
  • 云端存储
  • 提供各种工具
  • 可以和其他工具配合使用，比如下面的 pytorch, HF transformers, tensorflow, keras 等等
    
• 可以在里面使用 matplotlib
• 貌似是 tensorboard 的上位替代

注册账号

• 首先我们需要去官网注册账号，貌似不能使用vpn
  注册号后，按照教程创建一个团队，然后来到这个界面
  可以按照这个 Quickstart 的样例走一下。选择 Track Runs，接下来可以选择使用哪个工具训练的模型
  然后需要 pip install wandb 导包，以及 wandb login 登录
  

使用 HF Trainer + wandb 训练

• 我们调用官方给的样例
  我们发现其实新添了这几个内容：
  WANDB_PROJECT 环境变量：项目名
  WANDB_LOG_MODEL 环境变量：是否保存中继到wandb
  WANDB_WATCH环境变量
• 在 TrainingArguments 中，设置了 report_to="wandb"
  最后调用 wandb.finish() ，整体变化不大

# This script needs these libraries to be installed: 
#   numpy, transformers, datasets

import wandb 

import os
import numpy as np
from datasets import load_dataset
from transformers import TrainingArguments, Trainer
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 设置GPU编号
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return {"accuracy": np.mean(predictions == labels)}

print("Loading Dataset")
# download prepare the data
dataset = load_dataset("yelp_review_full")

print("Loading Tokenizer")

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

small_train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = dataset["test"].shuffle(seed=42).select(range(300))

small_train_dataset = small_train_dataset.map(tokenize_function, batched=True)
small_eval_dataset = small_train_dataset.map(tokenize_function, batched=True)

print("Loading Model")

# download the model
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=5)

# set the wandb project where this run will be logged
os.environ["WANDB_PROJECT"]="my-awesome-project"

# save your trained model checkpoint to wandb
os.environ["WANDB_LOG_MODEL"]="true"

# turn off watch to log faster
os.environ["WANDB_WATCH"]="false"

# pass "wandb" to the 'report_to' parameter to turn on wandb logging
training_args = TrainingArguments(
    output_dir='models',
    report_to="wandb",
    logging_steps=5, 
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="steps",
    eval_steps=20,
    max_steps = 100,
    save_steps = 100
)

print("Loading Trainer")

# define the trainer and start training
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

print("Training")

trainer.train()

# [optional] finish the wandb run, necessary in notebooks
wandb.finish()

• 在 wandb 网站中
  我们可以打开该 project。每一次运行相当于一次 run，我这里跑了三次所以就有三条线。
  这里主要是看 eval 验证集和 train 训练集的一些参数。
  
  
• 我们可以删掉不关心的面板，或者增添一个想看的面板
  但如果两个参数的值域变化比较大的话，在一个图里面比较难看清，所以比较相关的参数才建议放在一个图里。
  

低级 API

• 这上面是封装比较高级的 API，一般我们也都配合 transformers 库去用
  如果想用比较原生的 API，一般用法如下：
  首先调用 wandb.init() 方法
  然后使用 wandb.log(dict) 输出你要可视化的参数即可。

# train.py
import wandb
import random  # for demo script

wandb.login()

epochs = 10
lr = 0.01

run = wandb.init(
    # Set the project where this run will be logged
    project="my-awesome-project",
    # Track hyperparameters and run metadata
    config={
        "learning_rate": lr,
        "epochs": epochs,
    },
)

offset = random.random() / 5
print(f"lr: {lr}")

# simulating a training run
for epoch in range(2, epochs):
    acc = 1 - 2**-epoch - random.random() / epoch - offset
    loss = 2**-epoch + random.random() / epoch + offset
    print(f"epoch={epoch}, accuracy={acc}, loss={loss}")
    wandb.log({"accuracy": acc, "loss": loss})

# run.log_code()