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

这篇具有很好参考价值的文章主要介绍了【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()