chore: import upstream snapshot with attribution

python/ray/tune/tests/tutorial.py

@@ -0,0 +1,227 @@
+# ruff: noqa
+# isort: skip_file
+# Original Code: https://github.com/pytorch/examples/blob/master/mnist/main.py
+
+# fmt: off
+# __tutorial_imports_begin__
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+from ray import tune
+from ray.tune.schedulers import ASHAScheduler
+# __tutorial_imports_end__
+# fmt: on
+
+
+# fmt: off
+# __model_def_begin__
+class ConvNet(nn.Module):
+    def __init__(self):
+        super(ConvNet, self).__init__()
+        # In this example, we don't change the model architecture
+        # due to simplicity.
+        self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
+        self.fc = nn.Linear(192, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 3))
+        x = x.view(-1, 192)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=1)
+# __model_def_end__
+# fmt: on
+
+# fmt: off
+# __train_def_begin__
+
+# Change these values if you want the training to run quicker or slower.
+EPOCH_SIZE = 512
+TEST_SIZE = 256
+
+def train_func(model, optimizer, train_loader):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        # We set this just for the example to run quickly.
+        if batch_idx * len(data) > EPOCH_SIZE:
+            return
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+
+
+def test_func(model, data_loader):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (data, target) in enumerate(data_loader):
+            # We set this just for the example to run quickly.
+            if batch_idx * len(data) > TEST_SIZE:
+                break
+            data, target = data.to(device), target.to(device)
+            outputs = model(data)
+            _, predicted = torch.max(outputs.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    return correct / total
+# __train_def_end__
+
+
+# __train_func_begin__
+import os
+import tempfile
+
+from ray.tune import Checkpoint
+
+def train_mnist(config):
+    # Data Setup
+    mnist_transforms = transforms.Compose(
+        [transforms.ToTensor(),
+         transforms.Normalize((0.1307, ), (0.3081, ))])
+
+    train_loader = DataLoader(
+        datasets.MNIST("~/data", train=True, download=True, transform=mnist_transforms),
+        batch_size=64,
+        shuffle=True)
+    test_loader = DataLoader(
+        datasets.MNIST("~/data", train=False, transform=mnist_transforms),
+        batch_size=64,
+        shuffle=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = ConvNet()
+    model.to(device)
+
+    optimizer = optim.SGD(
+        model.parameters(), lr=config["lr"], momentum=config["momentum"])
+    for i in range(10):
+        train_func(model, optimizer, train_loader)
+        acc = test_func(model, test_loader)
+
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+            if (i + 1) % 5 == 0:
+                # This saves the model to the trial directory
+                torch.save(
+                    model.state_dict(),
+                    os.path.join(temp_checkpoint_dir, "model.pth")
+                )
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+            # Send the current training result back to Tune
+            tune.report({"mean_accuracy": acc}, checkpoint=checkpoint)
+
+# __train_func_end__
+# fmt: on
+
+# __eval_func_begin__
+search_space = {
+    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
+    "momentum": tune.uniform(0.1, 0.9),
+}
+
+# Uncomment this to enable distributed execution
+# `ray.init(address="auto")`
+
+# Download the dataset first
+datasets.MNIST("~/data", train=True, download=True)
+
+tuner = tune.Tuner(
+    train_mnist,
+    param_space=search_space,
+)
+results = tuner.fit()
+# __eval_func_end__
+
+# __plot_begin__
+dfs = {result.path: result.metrics_dataframe for result in results}
+[d.mean_accuracy.plot() for d in dfs.values()]
+# __plot_end__
+
+# __run_scheduler_begin__
+tuner = tune.Tuner(
+    train_mnist,
+    tune_config=tune.TuneConfig(
+        num_samples=20,
+        scheduler=ASHAScheduler(metric="mean_accuracy", mode="max"),
+    ),
+    param_space=search_space,
+)
+results = tuner.fit()
+
+# Obtain a trial dataframe from all run trials of this `tune.run` call.
+dfs = {result.path: result.metrics_dataframe for result in results}
+# __run_scheduler_end__
+
+# fmt: off
+# __plot_scheduler_begin__
+# Plot by epoch
+ax = None  # This plots everything on the same plot
+for d in dfs.values():
+    ax = d.mean_accuracy.plot(ax=ax, legend=False)
+# __plot_scheduler_end__
+# fmt: on
+
+# __run_searchalg_begin__
+from hyperopt import hp
+from ray.tune.search.hyperopt import HyperOptSearch
+
+space = {
+    "lr": hp.loguniform("lr", -10, -1),
+    "momentum": hp.uniform("momentum", 0.1, 0.9),
+}
+
+hyperopt_search = HyperOptSearch(space, metric="mean_accuracy", mode="max")
+
+tuner = tune.Tuner(
+    train_mnist,
+    tune_config=tune.TuneConfig(
+        num_samples=10,
+        search_alg=hyperopt_search,
+    ),
+)
+results = tuner.fit()
+
+# To enable GPUs, use this instead:
+# analysis = tune.run(
+#     train_mnist, config=search_space, resources_per_trial={'gpu': 1})
+
+# __run_searchalg_end__
+
+# __run_analysis_begin__
+best_result = results.get_best_result("mean_accuracy", mode="max")
+with best_result.checkpoint.as_directory() as checkpoint_dir:
+    state_dict = torch.load(os.path.join(checkpoint_dir, "model.pth"))
+
+model = ConvNet()
+model.load_state_dict(state_dict)
+# __run_analysis_end__
+
+from ray.tune.examples.mnist_pytorch_trainable import TrainMNIST
+
+# __trainable_run_begin__
+search_space = {
+    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
+    "momentum": tune.uniform(0.1, 0.9),
+}
+
+tuner = tune.Tuner(
+    TrainMNIST,
+    run_config=tune.RunConfig(stop={"training_iteration": 10}),
+    param_space=search_space,
+)
+results = tuner.fit()
+# __trainable_run_end__