Weights & Biases 机器学习实验跟踪工具详解

1. 项目简介

Weights & Biases（简称W&B）是一款强大的机器学习实验跟踪和可视化工具，由Weights & Biases公司开发。它提供了一套完整的工具链，用于跟踪实验指标、可视化训练过程、管理模型版本和团队协作，帮助机器学习从业者更高效地开发和部署模型。

1.1 主要功能

• 实验跟踪：记录实验参数、指标和输出
• 可视化：实时可视化训练过程和模型性能
• 模型版本管理：追踪模型版本和迭代历史
• 团队协作：共享实验结果和模型
• 超参数优化：自动搜索最佳超参数组合
• 模型监控：监控模型在生产环境中的性能

1.2 应用场景

• 机器学习模型开发和训练
• 深度学习模型调优
• 研究实验管理和复现
• 团队协作和知识共享
• 模型部署和监控

2. 安装与配置

2.1 安装方法

Weights & Biases可以通过pip安装：

# 安装Weights & Biases
pip install wandb

2.2 初始配置

安装后，需要初始化Weights & Biases并登录：

# 初始化Weights & Biases
wandb login

运行此命令后，会打开浏览器窗口，要求你登录或创建Weights & Biases账户。登录后，会自动完成配置。

2.3 项目设置

在使用Weights & Biases之前，需要创建一个项目：

1. 登录Weights & Biases网站（https://wandb.ai）
2. 点击"Create Project"按钮
3. 输入项目名称和描述
4. 选择项目可见性（公开或私有）
5. 点击"Create"完成创建

3. 核心概念

3.1 运行（Run）

运行是Weights & Biases中的基本单位，代表一次完整的实验执行。每个运行包含以下信息：

• 超参数和配置
• 训练和验证指标
• 模型权重和输出
• 代码版本和环境信息

3.2 项目（Project）

项目是运行的集合，用于组织和管理相关的实验。一个项目可以包含多个运行，方便进行比较和分析。

3.3 实验（Experiment）

实验是一组相关的运行，通常用于测试不同的超参数或模型架构。

3.4 指标（Metric）

指标是衡量模型性能的数值，如准确率、损失值等。Weights & Biases会自动跟踪和可视化这些指标。

3.5 超参数（Hyperparameter）

超参数是模型训练前设置的参数，如学习率、批量大小等。Weights & Biases会记录这些参数，方便比较不同配置的性能。

4. 基本使用

4.1 初始化W&B

在代码中初始化Weights & Biases：

import wandb

# 初始化W&B运行
wandb.init(
    project="my-project",  # 项目名称
    name="experiment-1",  # 运行名称
    config={  # 超参数和配置
        "learning_rate": 0.001,
        "batch_size": 32,
        "epochs": 10
    }
)

4.2 记录指标

在训练过程中记录指标：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型、损失函数和优化器
model = nn.Linear(10, 1)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=wandb.config.learning_rate)

# 训练循环
for epoch in range(wandb.config.epochs):
    # 训练代码...
    
    # 计算损失和准确率
    loss = criterion(outputs, targets)
    accuracy = calculate_accuracy(outputs, targets)
    
    # 记录指标
    wandb.log({
        "epoch": epoch,
        "loss": loss.item(),
        "accuracy": accuracy
    })

4.3 保存模型

保存模型到Weights & Biases：

# 保存模型
torch.save(model.state_dict(), "model.pth")

# 记录模型文件
wandb.save("model.pth")

4.4 可视化数据

可视化数据和模型预测：

# 可视化训练数据
wandb.log({
    "train_data": wandb.Image(train_data_sample)
})

# 可视化模型预测
wandb.log({
    "predictions": wandb.Table(
        data=[[x, y, y_pred] for x, y, y_pred in zip(features, targets, predictions)],
        columns=["feature", "target", "prediction"]
    )
})

5. 高级功能

5.1 超参数优化

使用Weights & Biases进行超参数优化：

import wandb
from wandb.sweep import SweepConfig

# 定义超参数搜索空间
config = {
    "method": "grid",  # 网格搜索
    "parameters": {
        "learning_rate": {
            "values": [0.001, 0.01, 0.1]
        },
        "batch_size": {
            "values": [16, 32, 64]
        },
        "optimizer": {
            "values": ["sgd", "adam"]
        }
    }
}

# 初始化 sweep
sweep_id = wandb.sweep(SweepConfig(config), project="my-project")

# 定义训练函数
def train():
    # 初始化W&B运行
    with wandb.init() as run:
        # 获取当前超参数
        config = run.config
        
        # 训练代码...
        
        # 记录指标
        wandb.log({"accuracy": accuracy, "loss": loss})

# 启动 sweep
wandb.agent(sweep_id, train)

5.2 模型版本管理

使用Weights & Biases管理模型版本：

# 保存模型并创建版本
model_artifact = wandb.Artifact(
    "model",
    type="model",
    description="My trained model",
    metadata={"accuracy": accuracy, "epoch": epoch}
)

# 添加模型文件
model_artifact.add_file("model.pth")

# 记录模型版本
run.log_artifact(model_artifact)

# 加载模型版本
artifact = run.use_artifact("my-project/model:latest")
artifact_dir = artifact.download()
model.load_state_dict(torch.load(f"{artifact_dir}/model.pth"))

5.3 团队协作

邀请团队成员协作：

1. 登录Weights & Biases网站
2. 进入项目页面
3. 点击"Settings"按钮
4. 选择"Members"选项
5. 输入团队成员的邮箱地址
6. 设置权限级别
7. 点击"Invite"发送邀请

5.4 模型监控

使用Weights & Biases监控模型在生产环境中的性能：

# 初始化监控运行
wandb.init(project="my-project", name="production-monitoring")

# 监控模型预测
while True:
    # 获取新数据
    data = get_new_data()
    
    # 模型预测
    predictions = model.predict(data)
    
    # 计算性能指标
    accuracy = calculate_accuracy(predictions, ground_truth)
    latency = calculate_latency()
    
    # 记录指标
    wandb.log({
        "accuracy": accuracy,
        "latency": latency,
        "predictions": wandb.Table(
            data=[[d, p] for d, p in zip(data, predictions)],
            columns=["data", "prediction"]
        )
    })
    
    # 等待新数据
    time.sleep(60)

6. 实用案例

6.1 深度学习模型训练

场景描述：使用Weights & Biases跟踪深度学习模型的训练过程。

实现步骤：

1. 初始化W&B运行
2. 定义模型和训练参数
3. 在训练循环中记录指标
4. 可视化训练过程
5. 保存和版本控制模型

代码示例：

import wandb
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 初始化W&B运行
wandb.init(
    project="mnist-classification",
    name="cnn-model",
    config={
        "learning_rate": 0.001,
        "batch_size": 64,
        "epochs": 10,
        "dropout": 0.2
    }
)

# 加载数据
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST("./data", train=False, transform=transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=wandb.config.batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=wandb.config.batch_size)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(wandb.config.dropout)
        self.dropout2 = nn.Dropout2d(wandb.config.dropout)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=wandb.config.learning_rate)

# 训练循环
for epoch in range(wandb.config.epochs):
    model.train()
    train_loss = 0
    train_correct = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
        pred = output.argmax(dim=1, keepdim=True)
        train_correct += pred.eq(target.view_as(pred)).sum().item()
    
    # 计算训练指标
    train_loss /= len(train_loader.dataset)
    train_accuracy = train_correct / len(train_loader.dataset)
    
    # 验证模型
    model.eval()
    test_loss = 0
    test_correct = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            test_correct += pred.eq(target.view_as(pred)).sum().item()
    
    # 计算验证指标
    test_loss /= len(test_loader.dataset)
    test_accuracy = test_correct / len(test_loader.dataset)
    
    # 记录指标
    wandb.log({
        "epoch": epoch,
        "train_loss": train_loss,
        "train_accuracy": train_accuracy,
        "test_loss": test_loss,
        "test_accuracy": test_accuracy
    })
    
    print(f"Epoch {epoch+1}: Train Loss: {train_loss:.4f}, Train Acc: {train_accuracy:.4f}, Test Loss: {test_loss:.4f}, Test Acc: {test_accuracy:.4f}")

# 保存模型
torch.save(model.state_dict(), "model.pth")
wandb.save("model.pth")

# 完成运行
wandb.finish()

6.2 超参数优化

场景描述：使用Weights & Biases进行超参数优化，找到最佳模型配置。

实现步骤：

1. 定义超参数搜索空间
2. 初始化sweep
3. 定义训练函数
4. 启动sweep进行超参数搜索
5. 分析结果，选择最佳超参数组合

代码示例：

import wandb
from wandb.sweep import SweepConfig
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义超参数搜索空间
config = {
    "method": "bayes",  # 贝叶斯优化
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform",
            "min": -6,  # 1e-6
            "max": -2   # 1e-2
        },
        "batch_size": {
            "values": [16, 32, 64]
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0.1,
            "max": 0.5
        },
        "optimizer": {
            "values": ["adam", "sgd"]
        }
    },
    "metric": {
        "name": "test_accuracy",
        "goal": "maximize"
    }
}

# 初始化 sweep
sweep_id = wandb.sweep(SweepConfig(config), project="mnist-hyperparameter-tuning")

# 定义模型
class Net(nn.Module):
    def __init__(self, dropout):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(dropout)
        self.dropout2 = nn.Dropout2d(dropout)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

# 定义训练函数
def train():
    # 初始化W&B运行
    with wandb.init() as run:
        # 获取当前超参数
        config = run.config
        
        # 加载数据
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        
        train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
        test_dataset = datasets.MNIST("./data", train=False, transform=transform)
        
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config.batch_size)
        
        # 初始化模型
        model = Net(dropout=config.dropout)
        
        # 选择优化器
        if config.optimizer == "adam":
            optimizer = optim.Adam(model.parameters(), lr=10**config.learning_rate)
        else:
            optimizer = optim.SGD(model.parameters(), lr=10**config.learning_rate, momentum=0.9)
        
        criterion = nn.CrossEntropyLoss()
        
        # 训练循环
        for epoch in range(5):  # 简化示例，只训练5个epoch
            model.train()
            train_loss = 0
            train_correct = 0
            
            for batch_idx, (data, target) in enumerate(train_loader):
                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()
                
                train_loss += loss.item()
                pred = output.argmax(dim=1, keepdim=True)
                train_correct += pred.eq(target.view_as(pred)).sum().item()
            
            # 计算训练指标
            train_loss /= len(train_loader.dataset)
            train_accuracy = train_correct / len(train_loader.dataset)
            
            # 验证模型
            model.eval()
            test_loss = 0
            test_correct = 0
            
            with torch.no_grad():
                for data, target in test_loader:
                    output = model(data)
                    test_loss += criterion(output, target).item()
                    pred = output.argmax(dim=1, keepdim=True)
                    test_correct += pred.eq(target.view_as(pred)).sum().item()
            
            # 计算验证指标
            test_loss /= len(test_loader.dataset)
            test_accuracy = test_correct / len(test_loader.dataset)
            
            # 记录指标
            run.log({
                "epoch": epoch,
                "train_loss": train_loss,
                "train_accuracy": train_accuracy,
                "test_loss": test_loss,
                "test_accuracy": test_accuracy
            })

# 启动 sweep
wandb.agent(sweep_id, train, count=20)  # 运行20次实验

6.3 模型版本管理和部署

场景描述：使用Weights & Biases管理模型版本并部署最佳模型。

实现步骤：

1. 训练多个模型并记录到W&B
2. 比较不同模型的性能
3. 选择最佳模型并创建版本
4. 部署模型到生产环境

代码示例：

import wandb
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 训练并记录多个模型
model_configs = [
    {"name": "model-1", "learning_rate": 0.001, "dropout": 0.2},
    {"name": "model-2", "learning_rate": 0.0005, "dropout": 0.3},
    {"name": "model-3", "learning_rate": 0.002, "dropout": 0.1}
]

for config in model_configs:
    # 初始化W&B运行
    run = wandb.init(
        project="model-versioning",
        name=config["name"],
        config=config
    )
    
    # 训练代码...
    # 省略训练代码，与前面示例类似
    
    # 假设我们已经训练了模型并计算了准确率
    test_accuracy = 0.98  # 示例值
    
    # 保存模型
    model_path = f"{config['name']}.pth"
    torch.save(model.state_dict(), model_path)
    
    # 创建模型 artifact
    model_artifact = wandb.Artifact(
        config["name"],
        type="model",
        description=f"Model with lr={config['learning_rate']}, dropout={config['dropout']}",
        metadata={"accuracy": test_accuracy}
    )
    
    # 添加模型文件
    model_artifact.add_file(model_path)
    
    # 记录模型版本
    run.log_artifact(model_artifact)
    
    # 完成运行
    run.finish()

# 加载最佳模型
api = wandb.Api()
project = api.project("model-versioning")

# 获取所有模型 artifacts
artifacts = project.artifacts(type="model")

# 找到准确率最高的模型
best_accuracy = 0
best_artifact = None

for artifact in artifacts:
    if "accuracy" in artifact.metadata:
        accuracy = artifact.metadata["accuracy"]
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_artifact = artifact

print(f"最佳模型: {best_artifact.name}, 准确率: {best_accuracy}")

# 下载最佳模型
best_artifact_dir = best_artifact.download()
best_model_path = f"{best_artifact_dir}/{best_artifact.name}.pth"

# 加载模型
model = Net(dropout=0.2)  # 假设使用默认dropout
model.load_state_dict(torch.load(best_model_path))
model.eval()

# 部署模型到生产环境
# 这里可以添加部署代码，如导出到ONNX、部署到API等
print("模型部署完成")

7. 总结与展望

Weights & Biases是一款功能强大的机器学习实验跟踪和可视化工具，为机器学习从业者提供了全面的实验管理解决方案。它的主要优势包括：

• 全面的实验跟踪：记录所有实验参数、指标和输出
• 强大的可视化能力：实时可视化训练过程和模型性能
• 模型版本管理：追踪模型版本和迭代历史
• 团队协作功能：方便团队共享实验结果和模型
• 超参数优化：自动搜索最佳超参数组合
• 模型监控：监控模型在生产环境中的性能

未来，Weights & Biases有望在以下方面继续发展：

• 更深入的模型分析工具：提供更详细的模型性能分析
• 更广泛的集成：与更多机器学习框架和工具集成
• 更强大的自动化功能：进一步自动化机器学习工作流
• 更丰富的可视化选项：提供更多类型的可视化图表
• 更完善的模型部署工具：简化模型从开发到部署的流程

通过使用Weights & Biases，机器学习从业者可以更高效地管理实验、优化模型、协作开发，并最终构建更好的机器学习系统。Weights & Biases的出现为机器学习工作流带来了显著的改进，成为现代机器学习开发中不可或缺的工具之一。