diff --git a/pkgs/development/python-modules/torch/tests/default.nix b/pkgs/development/python-modules/torch/tests/default.nix index e3f2ca44ba5a..90fb21d018d7 100644 --- a/pkgs/development/python-modules/torch/tests/default.nix +++ b/pkgs/development/python-modules/torch/tests/default.nix @@ -28,4 +28,6 @@ rec { feature = "rocm"; libraries = ps: [ ps.torchWithRocm ]; }; + + mnist-example = callPackage ./mnist-example { }; } diff --git a/pkgs/development/python-modules/torch/tests/mnist-example/default.nix b/pkgs/development/python-modules/torch/tests/mnist-example/default.nix new file mode 100644 index 000000000000..bb8b90aae5a9 --- /dev/null +++ b/pkgs/development/python-modules/torch/tests/mnist-example/default.nix @@ -0,0 +1,46 @@ +{ + lib, + linkFarm, + fetchurl, + writers, + torch, + torchvision, + runCommand, +}: +let + fashionMnistDataset = linkFarm "fashion-mnist-dataset" ( + lib.mapAttrsToList + (name: hash: { + inherit name; + path = fetchurl { + url = "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/${name}"; + inherit hash; + }; + }) + { + "train-images-idx3-ubyte.gz" = "sha256-Ou3jjWGGOQiteGE/ajLtJxYm3RKAC6JjZWlRI2kmioQ="; + "train-labels-idx1-ubyte.gz" = "sha256-oE8XE0rANWCkfjdk4RuS/JfeTRv6+LoaOqKa9UzJCEU="; + "t10k-images-idx3-ubyte.gz" = "sha256-NG5VuUjZc6l+WNI1Hd4WpIS9QV1FlSl2M7sI8D22oHM="; + "t10k-labels-idx1-ubyte.gz" = "sha256-Z9oXx26v/KVEbDNhqqtcPNbRwmCHZNNd+xhQsIa/jdU="; + } + ); + + mnist-script = writers.writePython3 "test_mnist" { + libraries = [ + torch + torchvision + ]; + flakeIgnore = [ "E501" ]; + } (builtins.readFile ./script.py); +in +runCommand "mnist" { } '' + mkdir -p data/FashionMNIST/raw + + for archive in `ls ${fashionMnistDataset}`; do + gzip -d < "${fashionMnistDataset}/$archive" > data/FashionMNIST/raw/"''${archive%.*}" + done + + ${mnist-script} + + touch $out +'' diff --git a/pkgs/development/python-modules/torch/tests/mnist-example/script.py b/pkgs/development/python-modules/torch/tests/mnist-example/script.py new file mode 100644 index 000000000000..4800125917bf --- /dev/null +++ b/pkgs/development/python-modules/torch/tests/mnist-example/script.py @@ -0,0 +1,160 @@ +import torch +from torch import nn +from torch.utils.data import DataLoader +from torchvision import datasets +from torchvision.transforms import ToTensor + +# Download training data from open datasets. +training_data = datasets.FashionMNIST( + root="data", + train=True, + download=False, + transform=ToTensor(), +) + +# Download test data from open datasets. +test_data = datasets.FashionMNIST( + root="data", + train=False, + download=False, + 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) + +for X, y in test_dataloader: + print(f"Shape of X [N, C, H, W]: {X.shape}") + print(f"Shape of y: {y.shape} {y.dtype}") + break + +device = ( + torch.accelerator.current_accelerator().type + if torch.accelerator.is_available() + else "cpu" +) +print(f"Using {device} device") + + +# Define model +class NeuralNetwork(nn.Module): + def __init__(self): + super().__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) + +loss_fn = nn.CrossEntropyLoss() +optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + + +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), + ) + + # Compute prediction error + pred = model(X) + loss = loss_fn(pred, y) + + # Backpropagation + loss.backward() + optimizer.step() + optimizer.zero_grad() + + if batch % 100 == 0: + loss, current = ( + loss.item(), + (batch + 1) * len(X), + ) + print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]") + + +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 + print( + f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n" + ) + + +epochs = 5 +for t in range(epochs): + print(f"Epoch {t + 1}\n-------------------------------") + train( + train_dataloader, + model, + loss_fn, + optimizer, + ) + test(test_dataloader, model, loss_fn) +print("Done!") + +torch.save(model.state_dict(), "model.pth") +print("Saved PyTorch Model State to model.pth") + +model = NeuralNetwork().to(device) +model.load_state_dict(torch.load("model.pth", weights_only=True)) + +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(): + x = x.to(device) + pred = model(x) + predicted, actual = ( + classes[pred[0].argmax(0)], + classes[y], + ) + print(f'Predicted: "{predicted}", Actual: "{actual}"')