python312Packages.torch: add mnist-example test (#404484)

pkgs/development/python-modules/torch/tests/default.nix

@@ -28,4 +28,6 @@ rec {
     feature = "rocm";
     libraries = ps: [ ps.torchWithRocm ];
   };
+
+  mnist-example = callPackage ./mnist-example { };
 }

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
+''

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}"')