训练应该如何与RayTune这样的超参数优化工具一起使用？

pip install ray[tune]

import torch
import torch.nn as nn
import torch.optim as optim
import ray
from ray import tune
from ray.tune.schedulers import ASHAScheduler

# 定义一个简单的神经网络
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 定义训练函数
def train_model(config):
    # 使用传入的超参数
    input_size = config["input_size"]
    hidden_size = config["hidden_size"]
    output_size = config["output_size"]
    lr = config["lr"]
    batch_size = config["batch_size"]
    epochs = config["epochs"]

    # 创建数据集和数据加载器
    train_data = torch.randn(1000, input_size)
    train_labels = torch.randint(0, output_size, (1000,))
    train_loader = torch.utils.data.DataLoader(
        dataset=list(zip(train_data, train_labels)),
        batch_size=batch_size,
        shuffle=True
    )

    # 创建模型、损失函数和优化器
    model = SimpleModel(input_size, hidden_size, output_size)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # 训练模型
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        # 计算平均损失
        avg_loss = running_loss / len(train_loader)
        
        # 使用 Ray Tune 报告结果
        tune.report(loss=avg_loss)

# 定义搜索空间
search_space = {
    "input_size": 20,
    "hidden_size": tune.choice([32, 64, 128]),
    "output_size": 10,
    "lr": tune.loguniform(1e-4, 1e-1),
    "batch_size": tune.choice([16, 32, 64]),
    "epochs": 10
}

# 使用 ASHA 调度器进行早停
scheduler = ASHAScheduler(
    metric="loss",
    mode="min",
    max_t=10,
    grace_period=1,
    reduction_factor=2
)

# 运行超参数优化
analysis = tune.run(
    train_model,
    config=search_space,
    num_samples=10,
    scheduler=scheduler
)

# 打印最佳配置
print("Best config: ", analysis.best_config)