将PyTorch代码无缝切换至Ray AIR
如果已经为某机器学习或数据分析编写了PyTorch代码，那么不必从头开始编写Ray AIR代码。相反，可以继续使用现有的代码，并根据需要逐步添加Ray AIR组件。
使用Ray AIR与现有的PyTorch训练代码，具有以下好处：
- 轻松在集群上进行分布式数据并行训练
- 自动检查点/容错和结果跟踪
- 并行数据预处理
- 与超参数调优的无缝集成
- 可扩展的批量预测
- 可扩展的模型服务
本教程将展示如何从现有的PyTorch训练代码开始使用Ray AIR，并将学习如何分布式训练和进行可扩展的批量预测。
代码示例
示例代码是PyTorch快速入门教程的代码。此代码在FashionMNIST数据集上训练一个神经网络分类器。 代码环境为《机器学习框架Ray -- 1.4 Ray RLlib的基本使用》中Anaconda创建的RayRLlib。
PyTorch代码编写
先编写PyTorch代码，后续再转换为Ray AIR代码。
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor下载Fashion MNIST数据集
# Download training data from open datasets.
training_data = datasets.FashionMNIST(
root="data",
train=True,
download=True,
transform=ToTensor(),
)
# Download test data from open datasets.
test_data = datasets.FashionMNIST(
root="data",
train=False,
download=True,
transform=ToTensor(),
)定义数据加载器
batch_size = 64
# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)定义并实例化分类器神经网络（PyTorch实现）
# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
# Define model
class NeuralNetwork(nn.Module):
def __init__(self):
super(NeuralNetwork, self).__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28*28, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, 10)
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
model = NeuralNetwork().to(device)
print(model)定义交叉熵与SGD优化器。
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)定义训练循环。
def train_epoch(dataloader, model, loss_fn, optimizer):
# 将函数从"train"改名为"train_epoch"，以避免与后面的 Ray Train 模块产生冲突
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % 100 == 0:
loss, current = loss.item(), batch * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")验证循环：
def test_epoch(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
return test_loss使用PyTorch进行训练！
epochs = 50
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train_epoch(train_dataloader, model, loss_fn, optimizer)
test_epoch(test_dataloader, model, loss_fn)
print("Done!")训练完成，保存.pth文件。
torch.save(model.state_dict(), "PyTorch2RayAIR_model.pth")
print("Saved PyTorch Model State to model.pth")提供一种将上述数据加载、优化器定义、损失函数定义、神经网络训练封装在一起的函数。
def train_func():
batch_size = 64
lr = 1e-3
epochs = 50
# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)
# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
model = NeuralNetwork().to(device)
print(model)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train_epoch(train_dataloader, model, loss_fn, optimizer)
test_epoch(test_dataloader, model, loss_fn)
print("Done!")使用上述封装函数train_func()进行训练，可获得同样效果但更简洁的逻辑。
train_func()将PyTorch代码迁移至Ray AIR
在上一小节PyTorch代码的基础上，为了在Ray AIR平台实现分布式训练，在多个workers之间分发训练。为此需要以下几个部分：通过对训练数据进行分片来使用数据并行训练、设置模型以在机器之间通信梯度更新、将结果报告回 Ray Train。
导入 Ray Train 和 Ray AIR Session：
import ray.train as train
from ray.air import session配置一个字典来配置超参数：
def train_func(config: dict):
batch_size = config["batch_size"]
lr = config["lr"]
epochs = config["epochs"]还需要以下几处PyTorch代码上的改动：
1. 需要根据工作器的数量动态调整工作器批处理大小
- batch_size_per_worker = batch_size // session.get_world_size()
2. 为分布式数据分片准备数据加载器
- train_dataloader = train.torch.prepare_data_loader(train_dataloader)
- test_dataloader = train.torch.prepare_data_loader(test_dataloader)
3. 为分布式梯度更新准备模型
- model = train.torch.prepare_model(model)
4. 捕获验证损失并将其报告给 Ray Train
- test_loss = test(test_dataloader, model, loss_fn)
- session.report(dict(loss=test_loss))
5. 在 train_epoch() 和 test_epoch() 函数中，将数据集大小除以workers进程数大小，确保在分布式训练中，每个设备处理的样本数量大致相同
- size = len(dataloader.dataset) // session.get_world_size()
6. 在 train_epoch() 函数中，可以摆脱设备映射。Ray Train 会为我们处理这个问题。
以上，只需要不到 10 行 Ray Train 特定代码，就可以在原始PyTorch代码的基础上，实现分布式训练的Ray AIR代码。
根据上述修改点，修改train_epoch() 函数：
def train_epoch(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) // session.get_world_size() # Divide by word size
model.train()
for batch, (X, y) in enumerate(dataloader):
# We don't need this anymore! Ray Train does this automatically:
# X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % 100 == 0:
loss, current = loss.item(), batch * len(X)
# print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")继续修改test_epoch() 函数。
def test_epoch(dataloader, model, loss_fn):
size = len(dataloader.dataset) // session.get_world_size() # Divide by word size
num_batches = len(dataloader)
model.eval()
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model(X)
test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
# print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
return test_loss最后，修改train_func()函数
import ray.train as train
from ray.air import session
def train_func(config: dict):
batch_size = config["batch_size"]
lr = config["lr"]
epochs = config["epochs"]
batch_size_per_worker = batch_size // session.get_world_size()
# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size_per_worker)
test_dataloader = DataLoader(test_data, batch_size=batch_size_per_worker)
train_dataloader = train.torch.prepare_data_loader(train_dataloader)
test_dataloader = train.torch.prepare_data_loader(test_dataloader)
model = NeuralNetwork()
model = train.torch.prepare_model(model)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
for t in range(epochs):
train_epoch(train_dataloader, model, loss_fn, optimizer)
test_loss = test_epoch(test_dataloader, model, loss_fn)
session.report(dict(loss=test_loss))
print("Done!")使用 Ray Train 的 TorchTrainer 来开始训练。
代码中设置了超参数config。
在 scaling_config 中，还可以配置要使用多少并行工作器以及是否要启用 GPU 训练（见注释）。
from ray.train.torch import TorchTrainer
from ray.air.config import ScalingConfig
# 使用CPU计算
device = torch.device("cpu")
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config={"lr": 1e-3, "batch_size": 64, "epochs": 50},
scaling_config=ScalingConfig(num_workers=20, use_gpu=False), # num_workers=gpu数量
)
# 使用GPU计算
#device = torch.device("gpu")
# trainer = TorchTrainer(
# train_loop_per_worker=train_func,
# train_loop_config={"lr": 1e-3, "batch_size": 64, "epochs": 50},
# scaling_config=ScalingConfig(num_workers=1, use_gpu=True), # num_workers=gpu数量
# )
result = trainer.fit()
print(f"Last result: {result.metrics}")进一步改进
上一小节的代码实现了并行训练模型，可以对代码进行一些改进，以便充分利用系统。
首先，我们应该启用检查点以在训练后访问训练模型。此外，我们应该优化数据加载以在工作器内进行。
- 启用checkpoint
启用检查点非常简单，只需要将带有模型状态的 Checkpoint 对象传递给 session.report() API。
- 将数据加载器移动到训练函数中
可能已经注意到一个警告：Warning: The actor TrainTrainable is very large (52 MiB). Check that its definition is not implicitly capturing a large array or other object in scope. Tip: use ray.put() to put large objects in the Ray object store.
这是因为在训练函数外加载数据，然后 Ray 序列化它以使其可访问于远程任务（可能在远程节点上执行！）。虽然只有 52 MB 的数据，但这还好。
但是如果这是完整的图像数据集，肯定不希望不必要地在集群中传输数据。相反，应该将数据集加载部分移动到 train_func() 中。这将在每台机器上下载数据到磁盘，并实现更高效的数据加载。
结果如下：
from ray.air import Checkpoint
def load_data():
# Download training data from open datasets.
training_data = datasets.FashionMNIST(
root="data",
train=True,
download=True,
transform=ToTensor(),
)
# Download test data from open datasets.
test_data = datasets.FashionMNIST(
root="data",
train=False,
download=True,
transform=ToTensor(),
)
return training_data, test_data
def train_func(config: dict):
batch_size = config["batch_size"]
lr = config["lr"]
epochs = config["epochs"]
batch_size_per_worker = batch_size // session.get_world_size()
training_data, test_data = load_data() # <- this is new!
# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size_per_worker)
test_dataloader = DataLoader(test_data, batch_size=batch_size_per_worker)
train_dataloader = train.torch.prepare_data_loader(train_dataloader)
test_dataloader = train.torch.prepare_data_loader(test_dataloader)
model = NeuralNetwork()
model = train.torch.prepare_model(model)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
for t in range(epochs):
train_epoch(train_dataloader, model, loss_fn, optimizer)
test_loss = test_epoch(test_dataloader, model, loss_fn)
checkpoint = Checkpoint.from_dict(
dict(epoch=t, model=model.state_dict())
)
session.report(dict(loss=test_loss), checkpoint=checkpoint)
print("Done!")再次训练：
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config={"lr": 1e-3, "batch_size": 64, "epochs": 50},
scaling_config=ScalingConfig(num_workers=20, use_gpu=False),
)
result = trainer.fit()训练时，Ray会显示一些训练有关的信息。
- Trial name: 试验的名称，可以用于标识试验。
- _time_this_iter_s: 当前迭代所花费的时间，以秒为单位。
- _timestamp: 试验的UNIX时间戳。
- _training_iteration: 训练迭代次数。
- date: 试验的日期和时间。
- done: 布尔值，表示试验是否已完成。
- episodes_total: （在强化学习中使用）完成的总情节数。
- experiment_id: 试验的唯一标识符。
- hostname: 运行试验的机器的主机名。
- iterations_since_restore: 自恢复以来的迭代次数。
- loss: 模型的损失值。
- node_ip: 运行试验的机器的IP地址。
- pid: 运行试验的进程ID。
- should_checkpoint: 一个布尔值，表示是否应该创建检查点。
- time_since_restore: 自恢复以来所花费的时间，以秒为单位。
- time_this_iter_s: 当前迭代所花费的时间，以秒为单位（与_time_this_iter_s相同）。
- time_total_s: 试验的总运行时间，以秒为单位。
- timestamp: 试验的UNIX时间戳（与_timestamp相同）。
- timesteps_since_restore: 自恢复以来的时间步数。
- timesteps_total: 完成的总时间步数。
- training_iteration: 训练迭代次数（与_training_iteration相同）。
- trial_id: 试验的唯一ID。
- warmup_time: 预热时间，以秒为单位。
训练完成，可以通过下列代码查看结果：
print(f"Last result: {result.metrics}")
print(f"Checkpoint: {result.checkpoint}")模型推理
将训练好的模型进行加载，并测试推理准确性。
定义Fashion Mnist数据集的10种分类。
def predict_from_model(model):
classes = [
"T-shirt/top",
"Trouser",
"Pullover",
"Dress",
"Coat",
"Sandal",
"Shirt",
"Sneaker",
"Bag",
"Ankle boot",
]
model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
pred = model(x)
predicted, actual = classes[pred[0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')可以使用保存的模型和现有的代码进行预测
from ray.train.torch import TorchCheckpoint
model = TorchCheckpoint.from_checkpoint(result.checkpoint).get_model(NeuralNetwork())
predict_from_model(model)使用下列函数循环预测
classes = [
"T-shirt/top",
"Trouser",
"Pullover",
"Dress",
"Coat",
"Sandal",
"Shirt",
"Sneaker",
"Bag",
"Ankle boot",
]
def predict_from_model(model, data):
model.eval()
with torch.no_grad():
for x, y in data:
pred = model(x)
predicted, actual = classes[pred[0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')
predict_from_model(model, [test_data[i] for i in range(10)])使用Ray AIR进行可扩展的批量预测
可以使用Ray AIR的BatchPredictor类进行可扩展的预测。
"Scalable batch prediction" 是指可以扩展的批处理预测，即可以处理大量数据的预测任务。通常，批处理预测需要处理一批数据，因此需要一种方法来处理大量数据。这可以通过使用分布式计算和并行处理来实现，以加快处理速度并扩展处理能力。例如，使用 Ray AIR 的 BatchPredictor 类，可以在分布式环境中处理大规模的预测任务，从而实现可扩展的批处理预测。
from ray.train.batch_predictor import BatchPredictor
from ray.train.torch import TorchPredictor
batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, TorchPredictor, model=NeuralNetwork())Batch predictors是与Ray Datasets配合使用的，这里我们将测试数据集转换为Ray Dataset。请注意，这不是很高效，您可以查看我们的其他教程以了解如何更有效地生成Ray Dataset。
import ray.data
ds = ray.data.from_items([x.numpy() for x, y in test_data], parallelism=8)然后，使用两个workers触发预测：
results = batch_predictor.predict(ds, batch_size=32, min_scoring_workers=2)results是另一个Ray Dataset。可以使用results.show()来查看我们的预测结果。
results.show()预测结果示例：
{'predictions': array([-3.6989684, -5.162467 , -4.1274934, -3.8681617, -4.6273932,
6.6691933, -3.4286706, 7.2931404, 2.5856407, 8.199261 ],
dtype=float32)}
{'predictions': array([ 4.6849933 , 0.62037057, 12.630945 , 2.9539034 ,
7.839991 , -6.660334 , 10.108643 , -19.705801 ,
3.2825718 , -14.7644 ], dtype=float32)}
...
{'predictions': array([ 12.824979 , 3.3830273, 6.0827456, 6.7726793, 1.9203707,
-10.963151 , 9.430065 , -16.93004 , 0.5109027, -12.396169 ],
dtype=float32)}
这个预测结果中，对于每一个可能的类别，都给出了一个置信度分数，这些分数被包含在一个长度为10的数组中。每个分数都是一个浮点数，代表了该样本属于对应类别的可能性大小。
在预测结果中，最大的分数是5.228492，出现在第9个位置，因此这个样本被预测为“Ankle boot”，同时它的置信度也相对较高。最小的分数是-4.9750695，出现在第2个位置，因此该样本被预测为“Trouser”，但是由于该分数非常低，因此可以认为该预测结果不是很可靠。
如果想将这些预测转换为类名（如原始示例中所示），可以使用map函数来完成：
predicted_classes = results.map_batches(
lambda batch: [classes[pred.argmax(0)] for pred in batch["predictions"]],
batch_size=32,
batch_format="pandas")为了查看的预测效果如何，将预测标签与一些实际标签放在一起进行比较：
real_classes = [classes[y] for x, y in test_data]
for predicted, real in zip(predicted_classes.take(), real_classes):
print((predicted, real))结果示例：
('Ankle boot', 'Ankle boot')
('Pullover', 'Pullover')
('Trouser', 'Trouser')
('Trouser', 'Trouser')
('Shirt', 'Shirt')
('Trouser', 'Trouser')
('Coat', 'Coat')
('Shirt', 'Shirt')
('Sandal', 'Sandal')
('Sneaker', 'Sneaker')
('Coat', 'Coat')
('Sandal', 'Sandal')
('Sandal', 'Sneaker')
('Dress', 'Dress')
('Coat', 'Coat')
('Trouser', 'Trouser')
('Pullover', 'Pullover')
('Shirt', 'Coat')
('Bag', 'Bag')
('T-shirt/top', 'T-shirt/top')在本训练模型中，20个测试案例只有1个被错误预测为'Shirt'的'Coat'，其他19个图像均分类正确，相对来讲准确度较高。
值得注意的是，Ray官方教程中，训练epochs均为5，这个数值较小但是计算速度很快。本文的参数，包括神经网络规模、训练回合数均适当增加，实现了85%左右的预测准确性。
总结
本教程演示了如何将现有的PyTorch代码转换为可与Ray AIR一起使用的代码。学习了如何使用Ray Train抽象使分布式训练成为可能、通过Ray AIR保存和检索模型检查点、加载模型进行批量预测。