跳到主要内容

PyTorch 与Weights & Biases

在深度学习中，实验管理是一个至关重要的环节。随着模型复杂度的增加，跟踪实验、记录指标和可视化结果变得愈发重要。Weights & Biases（W&B）是一个强大的工具，可以帮助你轻松管理这些任务。本文将介绍如何将PyTorch与W&B结合使用，以提升你的深度学习工作流程。

什么是Weights & Biases？​

Weights & Biases（W&B）是一个用于机器学习实验跟踪、可视化和协作的平台。它可以帮助你记录实验的配置、超参数、指标、模型权重等信息，并提供丰富的可视化工具，帮助你更好地理解模型的性能。

安装与设置​

首先，你需要安装W&B库。可以通过以下命令安装：

bash
pip install wandb

安装完成后，你需要登录W&B账户。如果你还没有账户，可以在W&B官网注册一个。登录命令如下：

bash
wandb login

在PyTorch中使用W&B​

1. 初始化W&B​

在开始实验之前，你需要初始化W&B。可以通过以下代码完成：

python
import wandb

wandb.init(project="my-pytorch-project", config={
    "learning_rate": 0.01,
    "epochs": 10,
    "batch_size": 32
})

wandb.init 函数会创建一个新的实验，并将配置参数记录到W&B中。project 参数指定了项目的名称，config 参数用于记录实验的配置。

2. 记录指标​

在训练过程中，你可以使用 wandb.log 函数记录指标。例如，记录训练损失和验证准确率：

python
for epoch in range(10):
    train_loss = train_one_epoch(model, train_loader, optimizer)
    val_accuracy = validate(model, val_loader)
    
    wandb.log({
        "epoch": epoch,
        "train_loss": train_loss,
        "val_accuracy": val_accuracy
    })

wandb.log 函数会将指定的指标记录到W&B中，并在W&B的仪表板上实时显示。

3. 记录模型权重​

你还可以使用W&B记录模型的权重。以下代码展示了如何保存模型权重并上传到W&B：

python
torch.save(model.state_dict(), "model.pth")
wandb.save("model.pth")

wandb.save 函数会将文件上传到W&B，并可以在W&B的仪表板中查看和下载。

4. 可视化​

W&B提供了丰富的可视化工具，帮助你分析实验结果。你可以在W&B的仪表板中查看损失曲线、准确率曲线、混淆矩阵等。

实际案例​

假设你正在训练一个图像分类模型，使用CIFAR-10数据集。以下是一个完整的示例代码：

python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import wandb

# 初始化W&B
wandb.init(project="cifar10-classification", config={
    "learning_rate": 0.001,
    "epochs": 10,
    "batch_size": 64
})

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

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=wandb.config.batch_size, shuffle=True)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=wandb.config.batch_size, shuffle=False)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = Net()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=wandb.config.learning_rate)

# 训练模型
for epoch in range(wandb.config.epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            wandb.log({"epoch": epoch, "train_loss": running_loss / 100})
            running_loss = 0.0

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

wandb.log({"val_accuracy": correct / total})

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

总结​

通过将PyTorch与Weights & Biases结合使用，你可以轻松地跟踪实验、记录指标并可视化模型性能。W&B的强大功能可以帮助你更好地理解模型的训练过程，并优化模型的性能。

附加资源​

练习​

  1. 尝试在W&B中创建一个新的项目，并使用不同的超参数训练模型，比较不同配置下的模型性能。
  2. 使用W&B的可视化工具分析模型的训练曲线，找出模型的最佳配置。
  3. 将模型权重上传到W&B，并在W&B的仪表板中查看和下载模型权重。
提示

如果你在使用W&B时遇到问题，可以参考官方文档或加入W&B的社区论坛，获取帮助和支持。