Ray框架解析，轻松入门Python分布式机器学习

import numpy as np
import ray
from ray.util.sgd.torch import TorchTrainer
from ray.util.sgd.torch_utils import to_torch

# 初始化Ray环境
ray.init(address="auto")

# 定义一个简单的PyTorch模型
class LinearModel:
    def __init__(self):
        self.model = to_torch(torch.nn.Linear(1, 1))
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.01)
        self.criterion = to_torch(torch.nn.MSELoss())

    def forward(self, x):
        return self.model(x)

    def train_step(self, data_batch):
        x, y = data_batch
        self.optimizer.zero_grad()
        outputs = self.forward(x)
        loss = self.criterion(outputs, y)
        loss.backward()
        self.optimizer.step()
        return {"loss": loss.item()}

    def predict(self, x):
        return self.forward(x).detach().cpu().numpy()

# 生成一些用于训练的随机数据
num_items = 1000
features = np.random.rand(num_items, 1)
labels = 5 * features + np.random.rand(num_items, 1)

# 创建一个TorchTrainer实例，并开始训练
trainer = TorchTrainer(
    model_creator=LinearModel,
    data_creator=lambda: (features, labels),
    loss_creator=lambda t: t["loss"],
    config={"batch_size": 32, "epochs": 3},
)

result = trainer.train()

# 输出训练结果
print(f"Loss after training: {result['loss']}")