如何在PyTorch中使用Tune#
在本教程中，我们将向您展示如何将Tune集成到您的PyTorch训练工作流中。我们将遵循PyTorch文档中的此教程来训练一个CIFAR10图像分类器。
超参数调优可以极大地影响一个模型的表现，决定一个模型是普通还是高准确率。通常，简单的调整，比如选择不同的学习率或更改网络层的大小，会对模型性能产生显着影响。幸运的是，Tune使得探索这些最佳参数组合变得容易，并且与PyTorch配合良好。
正如您所看到的，我们只需要做一些轻微的修改。特别地，我们需要
将数据加载和训练封装在函数中，
使一些网络参数可配置，
添加检查点（可选），
并定义模型调优的搜索空间。
备注
要运行此示例，您需要安装以下内容：
$ pip install ray torch torchvision
设置 / 导入#
让我们从导入开始：
import numpy as np
import os
import tempfile
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from filelock import FileLock
from torch.utils.data import random_split
import torchvision
import torchvision.transforms as transforms
from typing import Dict
import ray
from ray import train, tune
from ray.train import Checkpoint
from ray.tune.schedulers import ASHAScheduler
2023-11-17 17:06:18,417 INFO util.py:154 -- Outdated packages:
ipywidgets==7.7.1 found, needs ipywidgets>=8
Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.
2023-11-17 17:06:18,665 INFO util.py:154 -- Outdated packages:
ipywidgets==7.7.1 found, needs ipywidgets>=8
Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.
大多数导入是为了构建PyTorch模型。只有最后三个导入是为了Ray Tune。
数据加载器#
我们将数据加载器封装在自己的函数中，并传递一个全局数据目录。这样，我们可以在不同的试验之间共享数据目录。
def load_data(data_dir="./data"):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
# 我们在这里添加了文件锁，因为多个工作进程会希望
# 下载数据，这可能会导致覆盖，因为
# DataLoader 不是线程安全的。
with FileLock(os.path.expanduser("~/.data.lock")):
trainset = torchvision.datasets.CIFAR10(
root=data_dir, train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(
root=data_dir, train=False, download=True, transform=transform)
return trainset, testset
def load_test_data():
# 加载假数据以进行快速冒烟测试。
trainset = torchvision.datasets.FakeData(
128, (3, 32, 32), num_classes=10, transform=transforms.ToTensor()
)
testset = torchvision.datasets.FakeData(
16, (3, 32, 32), num_classes=10, transform=transforms.ToTensor()
)
return trainset, testset
可配置神经网络#
我们只能调整那些可配置的参数。在这个例子中，我们可以指定全连接层的层大小：
class Net(nn.Module):
def __init__(self, l1=120, l2=84):
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, l1)
self.fc2 = nn.Linear(l1, l2)
self.fc3 = nn.Linear(l2, 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
训练函数#
现在变得有趣，因为我们对示例 来自 PyTorch 文档 进行了一些更改。
完整的代码示例如下：
def train_cifar(config):
net = Net(config["l1"], config["l2"])
device = "cpu"
if torch.cuda.is_available():
device = "cuda:0"
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9)
# 通过 `get_checkpoint()` API 加载现有检查点。
if train.get_checkpoint():
loaded_checkpoint = train.get_checkpoint()
with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:
model_state, optimizer_state = torch.load(
os.path.join(loaded_checkpoint_dir, "checkpoint.pt")
)
net.load_state_dict(model_state)
optimizer.load_state_dict(optimizer_state)
if config["smoke_test"]:
trainset, _ = load_test_data()
else:
trainset, _ = load_data()
test_abs = int(len(trainset) * 0.8)
train_subset, val_subset = random_split(
trainset, [test_abs, len(trainset) - test_abs])
trainloader = torch.utils.data.DataLoader(
train_subset,
batch_size=int(config["batch_size"]),
shuffle=True,
num_workers=0 if config["smoke_test"] else 8,
)
valloader = torch.utils.data.DataLoader(
val_subset,
batch_size=int(config["batch_size"]),
shuffle=True,
num_workers=0 if config["smoke_test"] else 8,
)
for epoch in range(10): # 多次遍历数据集
running_loss = 0.0
epoch_steps = 0
for i, data in enumerate(trainloader):
# 获取输入；数据是一个包含 [输入, 标签] 的列表。
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
# 将参数梯度归零
optimizer.zero_grad()
# 前进 + 后退 + 优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印统计数据
running_loss += loss.item()
epoch_steps += 1
if i % 2000 == 1999: # 每2000个小批次打印一次
print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1,
running_loss / epoch_steps))
running_loss = 0.0
# 验证损失
val_loss = 0.0
val_steps = 0
total = 0
correct = 0
for i, data in enumerate(valloader, 0):
with torch.no_grad():
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
outputs = net(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
loss = criterion(outputs, labels)
val_loss += loss.cpu().numpy()
val_steps += 1
# 在此处，我们保存了一个检查点。它会被自动注册。
# Ray Tune 可能会通过 ``get_checkpoint()`` 访问
# 在未来的迭代中。
# 注意，要保存像检查点这样的文件，仍然需要将其放在一个目录下。
# 建立一个检查点。
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
path = os.path.join(temp_checkpoint_dir, "checkpoint.pt")
torch.save(
(net.state_dict(), optimizer.state_dict()), path
)
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
train.report(
{"loss": (val_loss / val_steps), "accuracy": correct / total},
checkpoint=checkpoint,
)
print("Finished Training")
如您所见，大部分代码是直接从示例中修改而来的。
测试集准确性#
通常，机器学习模型的性能是在未用于训练该模型的保留测试集上进行测试的。我们还将其封装在一个函数中：
def test_best_model(best_result, smoke_test=False):
best_trained_model = Net(best_result.config["l1"], best_result.config["l2"])
device = "cuda:0" if torch.cuda.is_available() else "cpu"
best_trained_model.to(device)
checkpoint_path = os.path.join(best_result.checkpoint.to_directory(), "checkpoint.pt")
model_state, optimizer_state = torch.load(checkpoint_path)
best_trained_model.load_state_dict(model_state)
if smoke_test:
_, testset = load_test_data()
else:
_, testset = load_data()
testloader = torch.utils.data.DataLoader(
testset, batch_size=4, shuffle=False, num_workers=2
)
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = best_trained_model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print("Best trial test set accuracy: {}".format(correct / total))
正如你所看到的，该函数还期望一个 device 参数，因此我们可以在 GPU 上进行测试集验证。
配置搜索空间#
最后，我们需要定义 Tune 的搜索空间。以下是一个示例：
config = {
"l1": tune.sample_from(lambda _: 2**np.random.randint(2, 9)),
"l2": tune.sample_from(lambda _: 2**np.random.randint(2, 9)),
"lr": tune.loguniform(1e-4, 1e-1),
"batch_size": tune.choice([2, 4, 8, 16]),
}
tune.sample_from() 函数使得定义自定义采样方法以获取超参数成为可能。在这个例子中，l1 和 l2 参数应为 4 到 256 之间的 2 的幂，因此可以是 4、8、16、32、64、128 或 256。lr（学习率）应该在 0.0001 到 0.1 之间均匀采样。最后，批量大小是 2、4、8 和 16 之间的选择。
在每个试验中，Tune 现在将随机从这些搜索空间中采样一种参数组合。然后，它将并行训练多个模型，并找到这些模型中表现最好的一个。我们还使用 ASHAScheduler，它会提前终止表现不佳的试验。
您可以指定 CPU 的数量，然后这些 CPU 将可用于例如增加 PyTorch DataLoader 实例的 num_workers。在每个试验中，选定数量的 GPU 对 PyTorch 可见。试验无法访问未请求的 GPU，因此您不必担心两个试验使用相同资源集。
在这里，我们还可以指定分数 GPU，因此诸如 gpus_per_trial=0.5 是完全有效的。这些试验将相互共享 GPU。您只需确保模型仍然适合 GPU 内存。
训练模型后，我们将找到表现最好的模型并从检查点文件加载训练后的网络。然后，我们获得测试集的准确率，并通过打印报告一切。
完整的主函数如下所示：
# 将其设置为 True，以使用小型合成数据集运行冒烟测试。
SMOKE_TEST = False
def main(num_samples=10, max_num_epochs=10, gpus_per_trial=2, smoke_test=False):
config = {
"l1": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
"l2": tune.sample_from(lambda _: 2 ** np.random.randint(2, 9)),
"lr": tune.loguniform(1e-4, 1e-1),
"batch_size": tune.choice([2, 4, 8, 16]),
"smoke_test": smoke_test,
}
scheduler = ASHAScheduler(
max_t=max_num_epochs,
grace_period=1,
reduction_factor=2)
tuner = tune.Tuner(
tune.with_resources(
tune.with_parameters(train_cifar),
resources={"cpu": 2, "gpu": gpus_per_trial}
),
tune_config=tune.TuneConfig(
metric="loss",
mode="min",
scheduler=scheduler,
num_samples=num_samples,
),
param_space=config,
)
results = tuner.fit()
best_result = results.get_best_result("loss", "min")
print("Best trial config: {}".format(best_result.config))
print("Best trial final validation loss: {}".format(
best_result.metrics["loss"]))
print("Best trial final validation accuracy: {}".format(
best_result.metrics["accuracy"]))
test_best_model(best_result, smoke_test=smoke_test)
main(num_samples=2, max_num_epochs=2, gpus_per_trial=0, smoke_test=SMOKE_TEST)
如果你运行这段代码，一个示例输出可能如下所示：
试验次数：10 (10 终止)
+-------------------------+------------+-------+------+------+-------------+--------------+---------+------------+----------------------+
| 试验名称 | 状态 | loc | l1 | l2 | lr | batch_size | 损失 | 准确率 | 训练迭代 |
|-------------------------+------------+-------+------+------+-------------+--------------+---------+------------+----------------------|
| train_cifar_87d1f_00000 | 终止 | | 64 | 4 | 0.00011629 | 2 | 1.87273 | 0.244 | 2 |
| train_cifar_87d1f_00001 | 终止 | | 32 | 64 | 0.000339763 | 8 | 1.23603 | 0.567 | 8 |
| train_cifar_87d1f_00002 | 终止 | | 8 | 16 | 0.00276249 | 16 | 1.1815 | 0.5836 | 10 |
| train_cifar_87d1f_00003 | 终止 | | 4 | 64 | 0.000648721 | 4 | 1.31131 | 0.5224 | 8 |
| train_cifar_87d1f_00004 | 终止 | | 32 | 16 | 0.000340753 | 8 | 1.26454 | 0.5444 | 8 |
| train_cifar_87d1f_00005 | 终止 | | 8 | 4 | 0.000699775 | 8 | 1.99594 | 0.1983 | 2 |
| train_cifar_87d1f_00006 | 终止 | | 256 | 8 | 0.0839654 | 16 | 2.3119 | 0.0993 | 1 |
| train_cifar_87d1f_00007 | 终止 | | 16 | 128 | 0.0758154 | 16 | 2.33575 | 0.1327 | 1 |
| train_cifar_87d1f_00008 | 终止 | | 16 | 8 | 0.0763312 | 16 | 2.31129 | 0.1042 | 4 |
| train_cifar_87d1f_00009 | 终止 | | 128 | 16 | 0.000124903 | 4 | 2.26917 | 0.1945 | 1 |
+-------------------------+------------+-------+------+------+-------------+--------------+---------+------------+----------------------+
最佳试验配置：{'l1': 8, 'l2': 16, 'lr': 0.0027624906698231976, 'batch_size': 16, 'data_dir': '...'}
最佳试验最终验证损失：1.1815014744281769
最佳试验最终验证准确率：0.5836
最佳试验测试集准确率：0.5806
如你所见，大多数试验已早期停止，以避免浪费资源。 表现最佳的试验达到了约58%的验证准确率，这可以在测试集中得到验证。
所以就这样！你现在可以调整你的PyTorch模型的参数了。
查看更多PyTorch示例#
MNIST PyTorch 示例: 将PyTorch MNIST示例转换为使用Tune的函数式API。 还展示了如何轻松将依赖argparse的内容转换为使用Tune。
PBT ConvNet 示例: 示例使用函数API进行检查点训练的卷积网络。
MNIST PyTorch 可训练示例: 将PyTorch MNIST示例转换为使用Tune的可训练API。 还使用了HyperBand调度器，并在最后对模型进行了检查点存储。