:orphan: .. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "tutorials/hpo_nnictl/model.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_tutorials_hpo_nnictl_model.py: Port PyTorch Quickstart to NNI ============================== This is a modified version of `PyTorch quickstart`_. It can be run directly and will have the exact same result as original version. Furthermore, it enables the ability of auto tuning with an NNI *experiment*, which will be detailed later. It is recommended to run this script directly first to verify the environment. There are 2 key differences from the original version: 1. In `Get optimized hyperparameters`_ part, it receives generated hyperparameters. 2. In `Train model and report accuracy`_ part, it reports accuracy metrics to NNI. .. _PyTorch quickstart: https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html .. GENERATED FROM PYTHON SOURCE LINES 21-28 .. code-block:: default import nni import torch from torch import nn from torch.utils.data import DataLoader from torchvision import datasets from torchvision.transforms import ToTensor .. GENERATED FROM PYTHON SOURCE LINES 29-32 Hyperparameters to be tuned --------------------------- These are the hyperparameters that will be tuned. .. GENERATED FROM PYTHON SOURCE LINES 32-38 .. code-block:: default params = { 'features': 512, 'lr': 0.001, 'momentum': 0, } .. GENERATED FROM PYTHON SOURCE LINES 39-43 Get optimized hyperparameters ----------------------------- If run directly, :func:`nni.get_next_parameter` is a no-op and returns an empty dict. But with an NNI *experiment*, it will receive optimized hyperparameters from tuning algorithm. .. GENERATED FROM PYTHON SOURCE LINES 43-47 .. code-block:: default optimized_params = nni.get_next_parameter() params.update(optimized_params) print(params) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none {'features': 512, 'lr': 0.001, 'momentum': 0} .. GENERATED FROM PYTHON SOURCE LINES 48-50 Load dataset ------------ .. GENERATED FROM PYTHON SOURCE LINES 50-58 .. code-block:: default training_data = datasets.FashionMNIST(root="data", train=True, download=True, transform=ToTensor()) test_data = datasets.FashionMNIST(root="data", train=False, download=True, transform=ToTensor()) batch_size = 64 train_dataloader = DataLoader(training_data, batch_size=batch_size) test_dataloader = DataLoader(test_data, batch_size=batch_size) .. GENERATED FROM PYTHON SOURCE LINES 59-61 Build model with hyperparameters -------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 61-86 .. code-block:: default device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using {device} device") 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, params['features']), nn.ReLU(), nn.Linear(params['features'], params['features']), nn.ReLU(), nn.Linear(params['features'], 10) ) def forward(self, x): x = self.flatten(x) logits = self.linear_relu_stack(x) return logits model = NeuralNetwork().to(device) loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=params['lr'], momentum=params['momentum']) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none Using cpu device .. GENERATED FROM PYTHON SOURCE LINES 87-89 Define train and test --------------------- .. GENERATED FROM PYTHON SOURCE LINES 89-115 .. code-block:: default def train(dataloader, model, loss_fn, optimizer): size = len(dataloader.dataset) model.train() for batch, (X, y) in enumerate(dataloader): X, y = X.to(device), y.to(device) pred = model(X) loss = loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() def test(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 return correct .. GENERATED FROM PYTHON SOURCE LINES 116-119 Train model and report accuracy ------------------------------- Report accuracy metrics to NNI so the tuning algorithm can suggest better hyperparameters. .. GENERATED FROM PYTHON SOURCE LINES 119-126 .. code-block:: default epochs = 5 for t in range(epochs): print(f"Epoch {t+1}\n-------------------------------") train(train_dataloader, model, loss_fn, optimizer) accuracy = test(test_dataloader, model, loss_fn) nni.report_intermediate_result(accuracy) nni.report_final_result(accuracy) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none Epoch 1 ------------------------------- [2022-03-21 01:09:37] INFO (nni/MainThread) Intermediate result: 0.461 (Index 0) Epoch 2 ------------------------------- [2022-03-21 01:09:42] INFO (nni/MainThread) Intermediate result: 0.5529 (Index 1) Epoch 3 ------------------------------- [2022-03-21 01:09:47] INFO (nni/MainThread) Intermediate result: 0.6155 (Index 2) Epoch 4 ------------------------------- [2022-03-21 01:09:52] INFO (nni/MainThread) Intermediate result: 0.6345 (Index 3) Epoch 5 ------------------------------- [2022-03-21 01:09:56] INFO (nni/MainThread) Intermediate result: 0.6505 (Index 4) [2022-03-21 01:09:56] INFO (nni/MainThread) Final result: 0.6505 .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 24.441 seconds) .. _sphx_glr_download_tutorials_hpo_nnictl_model.py: .. only :: html .. container:: sphx-glr-footer :class: sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: model.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: model.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_