mnistvit 1.0.0__py3-none-any.whl

mnistvit/__init__.py

@@ -0,0 +1,6 @@
+"""
+`mnistvit` is a package providing a vision transformer for training on MNIST.
+
+The package includes functions for training and evaluation, as well as hyperparameter
+optimization.
+"""

mnistvit/__main__.py

@@ -0,0 +1,3 @@
+from .predict import main
+
+main()

mnistvit/model.py

@@ -0,0 +1,200 @@
+from math import prod
+from numbers import Number
+from typing import Dict, List, Literal
+
+from torch import Tensor, cat, nn, randn
+
+
+class VisionTransformer(nn.Module):
+    """Configurable vision transformer (ViT).
+
+    Implements the vision transformer as proposed by Dosovitskiy et al., ICLR 2021.
+
+    Args:
+        num_channels (int): Number of channels of the input.
+        input_sizes (list of int): Spatial sizes of the input.
+        output_size (int): Size of the output layer.
+        patch_size (int): Size of a patch in one dimension.
+        num_heads (int): Number of attention heads in each encoder block.
+        latent_size_multiplier (int): Yields the size of the embedding when multiplied
+            with `num_heads`.
+        num_layers (int): Number of encoder blocks.
+        encoder_size (int): Number of hidden units in each encoder MLP.
+        head_size (int or list of int): Sizes of hidden layers in MLP head.
+        dropout (float, optional): Dropout probabilities of embedding, encoder and MLP
+            head.  Default: 0.
+        encoder_activation (str, optional): Encoder activation function string, either
+            `'relu'` or `'gelu'`.  Default: `'gelu'`.
+        head_activation (str, optional): MLP head activation function string, `'relu'`,
+            `'gelu'` or `'tanh'`.  Default: `'gelu'`.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        input_sizes: List[int],
+        output_size: int,
+        patch_size: int,
+        num_heads: int,
+        latent_size_multiplier: int,
+        num_layers: int,
+        encoder_size: int,
+        head_size: int | List[int],
+        dropout: float = 0,
+        encoder_activation: Literal["relu", "gelu"] = "gelu",
+        head_activation: Literal["relu", "gelu", "tanh"] = "gelu",
+    ) -> None:
+        super().__init__()
+        self.kwargs = {
+            "num_channels": num_channels,
+            "input_sizes": input_sizes,
+            "output_size": output_size,
+            "patch_size": patch_size,
+            "num_heads": num_heads,
+            "latent_size_multiplier": latent_size_multiplier,
+            "num_layers": num_layers,
+            "encoder_size": encoder_size,
+            "head_size": head_size,
+            "dropout": dropout,
+            "encoder_activation": encoder_activation,
+            "head_activation": head_activation,
+        }
+        latent_size = latent_size_multiplier * num_heads
+        self.embedding = Embedding(
+            num_channels, input_sizes, patch_size, latent_size, dropout
+        )
+        layer_norm = nn.LayerNorm(latent_size)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=latent_size,
+            nhead=num_heads,
+            dim_feedforward=encoder_size,
+            dropout=dropout,
+            activation=encoder_activation,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers, layer_norm, enable_nested_tensor=False
+        )
+        if isinstance(head_size, Number):
+            head_size = [head_size]
+        self.mlp_head = MLP(
+            input_size=latent_size,
+            output_size=output_size,
+            hidden_sizes=head_size,
+            dropout=dropout,
+            activation=head_activation,
+        )
+
+    def forward(self, data: Tensor) -> Tensor:
+        data = self.embedding(data)
+        data = self.encoder(data)
+        # Take encoder output corresponding to class_token
+        output = self.mlp_head(data[:, 0, :])
+        return output
+
+
+class Embedding(nn.Module):
+    """An embedding for a vision transformer.
+
+    Splits the input into patches, and projects the patches.  Also adds a class token
+    and a position embedding.
+
+    Args:
+        num_channels (int): Number of channels of the input.
+        input_sizes (list of int): Spatial sizes of the input.
+        patch_size (int): Size of a patch in one dimension.
+        latent_size (int): Size of the embedding.
+        dropout (float, optional): Dropout probability.  Default: 0.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        input_sizes: List[int],
+        patch_size: int,
+        latent_size: int,
+        dropout: float = 0,
+    ) -> None:
+        super().__init__()
+        # Use Unfold to split the input image into patches
+        self.unfold = nn.Unfold(kernel_size=patch_size, stride=patch_size)
+        num_patches = prod(
+            [(input_size - patch_size) // patch_size + 1 for input_size in input_sizes]
+        )
+        flattened_size = num_channels * patch_size ** len(input_sizes)
+        self.linear = nn.Linear(flattened_size, latent_size)
+        self.class_token = nn.Parameter(randn(1, 1, latent_size))
+        self.position_embeddings = nn.Parameter(randn(1, num_patches + 1, latent_size))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, data: Tensor) -> Tensor:
+        # Split images into patches and flatten into
+        # batch_size x num_patches x (num_channels * patch_size ** len(input_sizes))
+        batch_size = data.shape[0]
+        data = self.unfold(data).permute(0, 2, 1).contiguous()
+        data = self.linear(data)
+        output = self.dropout(
+            cat([self.class_token.expand(batch_size, -1, -1), data], dim=1)
+            + self.position_embeddings
+        )
+        return output
+
+
+class MLP(nn.Module):
+    """Configurable multilayer perceptron.
+
+    Args:
+        input_size (int): Size of the input layer.
+        output_size (int): Size of the output layer.
+        hidden_sizes (list of int, optional): Sizes of hidden layers.  Default: `None`.
+        dropout (float or list of float, optional): Dropout probabilities of each
+            hidden layer.  If `None`, no dropout will be used.  If single float, the
+            same dropout probability will be used for all hidden layers.
+            Default: `None`.
+        activation (str, optional): Activation function string, `'relu'`, `'gelu'` or
+            `'tanh'`.  Default: `'relu'`.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        hidden_sizes: List[int] = None,
+        dropout: float | List[float] = None,
+        activation: Literal["relu", "gelu", "tanh"] = "relu",
+    ) -> None:
+        super().__init__()
+        self.flatten = nn.Flatten()
+        if hidden_sizes is None:
+            hidden_sizes = []
+        layers = [
+            nn.Linear(i, j)
+            for i, j in zip([input_size] + hidden_sizes, hidden_sizes + [output_size])
+        ]
+        if dropout is None:
+            dropout_modules = []
+        elif isinstance(dropout, Number):
+            dropout_modules = [nn.Dropout(dropout) for _ in range(len(hidden_sizes))]
+        else:
+            dropout_modules = [nn.Dropout(rate) for rate in dropout]
+        modules = []
+        for i, layer in enumerate(layers):
+            modules.append(layer)
+            if i < len(layers) - 1:
+                if activation == "relu":
+                    modules.append(nn.ReLU())
+                elif activation == "gelu":
+                    modules.append(nn.GELU())
+                elif activation == "tanh":
+                    modules.append(nn.Tanh())
+                else:
+                    raise ValueError(f"unknown activation '{activation}'")
+                if len(dropout_modules) > i:
+                    modules.append(dropout_modules[i])
+        self.linear_stack = nn.Sequential(*modules)
+
+    def forward(self, data: Tensor) -> Tensor:
+        data = self.flatten(data)
+        output = self.linear_stack(data)
+        return output

mnistvit/predict.py

@@ -0,0 +1,244 @@
+import argparse
+from typing import Dict
+
+import torch
+
+from .preprocess import get_test_loader_mnist, read_digit_image
+from .utils import FILE_LIKE, load_model
+
+
+def test_mnist(
+    config: Dict,
+    data_dir: str,
+    model_file: FILE_LIKE,
+    use_loss: bool = True,
+    use_accuracy: bool = True,
+    device: torch.device = "cpu",
+) -> None:
+    """Loads a model, tests it on MNIST and prints the results.
+
+    Args:
+        config (dict): Test configuration with `'batch_size'`.
+        data_dir (str): Directory of the MNIST dataset.
+        model_file (FILE_LIKE): File name to load the model from.
+        use_loss (bool, optional): If true, evaluates the loss on the test set.
+            Default: `True`.
+        use_accuracy (bool, optional): If true, evaluates the accuracy on the test set.
+            Default: `True`.
+        device (torch.device, optional): Device to evaluate the model on.
+            Default: `'cpu'`.
+    """
+    model = load_model(model_file, device)
+    test_loader = get_test_loader_mnist(data_dir, config["batch_size"])
+    if use_loss:
+        loss_fn = torch.nn.CrossEntropyLoss()
+        loss = prediction_loss(model, test_loader, loss_fn, device)
+        print("Test loss: ", loss)
+    if use_accuracy:
+        acc = prediction_accuracy(model, test_loader, device)
+        print("Test accuracy: ", acc)
+
+
+def predict_file(
+    image_file: FILE_LIKE, model_file: FILE_LIKE, device: torch.device = "cpu"
+) -> int:
+    """Loads a model and classifies a digit from an image file.
+
+    Args:
+        image_file (FILE_LIKE): The image file.
+        model_file (FILE_LIKE): File name to load the model from.
+        device (torch.device, optional): Device to evaluate the model on.
+            Default: `'cpu'`.
+
+    Returns:
+        int: Predicted class label.
+    """
+    model = load_model(model_file, device)
+    image = read_digit_image(image_file)
+    predicted = predict_single_image(model, image, device)
+    return predicted
+
+
+def prediction_loss(
+    model: torch.nn.Module,
+    data_loader: torch.utils.data.DataLoader,
+    loss_fn: torch.nn.Module,
+    device: torch.device = "cpu",
+) -> float:
+    """Evaluates the model loss on the data loader.
+
+    Args:
+        model (torch.nn.Module): Model to evaluate.
+        data_loader (torch.utils.data.DataLoader): Data loader for evaluation.
+        loss_fn (torch.nn.Module): Loss function for evaluation.
+        device (torch.device, optional): Device to evaluate the model on.
+            Default: `'cpu'`.
+
+    Returns:
+        float: Calculated loss.
+    """
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for data, target in data_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            loss += loss_fn(output, target).cpu().numpy()
+    loss /= len(data_loader)
+    return loss
+
+
+def prediction_accuracy(
+    model: torch.nn.Module,
+    data_loader: torch.utils.data.DataLoader,
+    device: torch.device = "cpu",
+) -> float:
+    """Evaluates the model accuracy on the data loader.
+
+    Args:
+        model (torch.nn.Module): Model to evaluate.
+        data_loader (torch.utils.data.DataLoader): Data loader for evaluation.
+        device (torch.device, optional): Device to evaluate the model on.
+            Default: `'cpu'`.
+
+    Returns:
+        float: Calculated accuracy.
+    """
+    correct = 0
+    total = 0
+    for data, target in data_loader:
+        data, target = data.to(device), target.to(device)
+        predicted = classify(model, data)
+        total += target.size(0)
+        correct += (predicted == target).sum().item()
+    return correct / total
+
+
+def predict_single_image(
+    model: torch.nn.Module, image: torch.FloatTensor, device: torch.device = "cpu"
+) -> int:
+    """Uses model to classify a digit image.
+
+    Args:
+        model (torch.nn.Module): Model to use for classification.
+        image (torch.FloatTensor): Preprocessed digit image.
+        device (torch.device, optional): Device to evaluate the model on.
+            Default: `'cpu'`.
+
+    Returns:
+        int: Predicted class label.
+    """
+    data = image.unsqueeze(0)  # Add batch dimension
+    data = data.to(device)
+    predicted = classify(model, data).cpu().item()
+    return predicted
+
+
+def classify(model: torch.nn.Module, data: torch.utils.data.Sampler) -> torch.Tensor:
+    """Uses model to classify given data.
+
+    Args:
+        model (torch.nn.Module): Model to use for classification.
+        data (torch.utils.data.Sampler): Data sampler.
+
+    Returns:
+        torch.Tensor: Predicted class labels.
+    """
+    output = eval_output(model, data)
+    predicted = torch.max(output.data, dim=1)[1]
+    return predicted
+
+
+def class_log_probs(
+    model: torch.nn.Module, data: torch.utils.data.Sampler
+) -> torch.Tensor:
+    """Evaluates model log probabilities of all classes on given data.
+
+    Args:
+        model (torch.nn.Module): Model to use for evaluation.
+        data (torch.utils.data.Sampler): Data sampler.
+
+    Returns:
+        torch.Tensor: Log probabilities of classes.
+    """
+    output = eval_output(model, data)
+    log_probs = torch.nn.functional.log_softmax(output.data, dim=1)
+    return log_probs
+
+
+def eval_output(model: torch.nn.Module, data: torch.utils.data.Sampler) -> torch.Tensor:
+    """Evaluates the output of a model on given data.
+
+    Args:
+        model (torch.nn.Module): Model to use for evaluation.
+        data (torch.utils.data.Sampler): Data sampler.
+
+    Returns:
+        torch.Tensor: Model output.
+    """
+    model.eval()
+    with torch.no_grad():
+        output = model(data)
+    return output
+
+
+def main() -> None:
+    """Processes command line arguments with prediction."""
+    parser = argparse.ArgumentParser(description="MNIST Prediction")
+    parser.add_argument(
+        "--image-file",
+        type=str,
+        default=None,
+        metavar="FILE",
+        help="image file to predict (default: None)",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        metavar="N",
+        help="input batch size for testing (default: 32)",
+    )
+    parser.add_argument(
+        "--model-file",
+        type=str,
+        default="model.pt",
+        metavar="FILE",
+        help="file to load the model from (default: 'model.pt')",
+    )
+    parser.add_argument(
+        "--use-loss",
+        action="store_true",
+        default=False,
+        help="enables test loss calculation",
+    )
+    parser.add_argument(
+        "--use-accuracy",
+        action="store_true",
+        default=False,
+        help="enables test accuracy calculation",
+    )
+    parser.add_argument(
+        "--no-cuda", action="store_true", default=False, help="disables CUDA testing"
+    )
+    args = parser.parse_args()
+    no_cuda = args.no_cuda or not torch.cuda.is_available()
+    device = torch.device("cpu" if no_cuda else "cuda")
+    config = {
+        "batch_size": args.batch_size,
+    }
+    if args.image_file is not None:
+        predicted = predict_file(args.image_file, args.model_file, device)
+        print(predicted)
+    test_mnist(
+        config,
+        data_dir="data",
+        model_file=args.model_file,
+        use_loss=args.use_loss,
+        use_accuracy=args.use_accuracy,
+        device=device,
+    )
+
+
+if __name__ == "__main__":
+    main()

mnistvit/preprocess.py

@@ -0,0 +1,124 @@
+from typing import Tuple
+
+import torch
+from torch.utils.data import DataLoader, random_split
+from torchvision import datasets, transforms
+from torchvision.io import ImageReadMode, read_image
+
+from .utils import FILE_LIKE
+
+
+def preprocess_mnist(
+    data_dir: str, train: bool, use_augmentation: bool = False
+) -> Tuple:
+    """Preprocesses the MNIST dataset.
+
+    Normalizes the MNIST dataset with training mean and variance and eventually
+    augments the dataset with random affine transformations.
+
+    Args:
+        data_dir (str): Directory of the MNIST dataset.
+        train (bool): If true, loads the training set, else the test set.
+        use_augmentation (bool, optional): If true, augments the dataset with random
+            affine transformations.  Default: `False`.
+
+    Returns:
+        tuple: Image, target pairs.
+    """
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+    if use_augmentation:
+        # Data augmentation
+        transform = transforms.Compose(
+            [
+                transforms.RandomAffine(
+                    degrees=16, translate=(0.1, 0.1), scale=(0.9, 1.1)
+                ),
+                transform,
+            ]
+        )
+    dataset = datasets.MNIST(data_dir, train=train, download=True, transform=transform)
+    return dataset
+
+
+def get_train_loaders_mnist(
+    data_dir: str,
+    batch_size: int,
+    train_fraction: float = 1.0,
+    use_augmentation: bool = True,
+) -> Tuple:
+    """Training loaders of the MNIST dataset.
+
+    Args:
+        data_dir (str): Directory of the MNIST dataset.
+        batch_size (int): Size of the batches of the training loaders.
+        train_fraction (float, optional): Fraction of the set used for the training
+            loader.  The remainder is used for the validation loader.  Default: 1.0.
+        use_augmentation (bool, optional): If true, augments the dataset with random
+            affine transformations.  Default: `True`.
+
+    Returns:
+        tuple: Training loader and validation loader, where the latter is `None` if
+            `train_fraction` is 1.
+    """
+    assert 0 <= train_fraction <= 1
+    loader_kwargs = {"shuffle": True, "batch_size": batch_size}
+    dataset = preprocess_mnist(data_dir, train=True, use_augmentation=use_augmentation)
+    if train_fraction == 1.0:
+        train_loader = DataLoader(dataset, **loader_kwargs)
+        val_loader = None
+    else:
+        num_train = int(len(dataset) * train_fraction)
+        train_set, val_set = random_split(
+            dataset, [num_train, len(dataset) - num_train]
+        )
+        train_loader = DataLoader(train_set, **loader_kwargs)
+        val_loader = DataLoader(val_set, **loader_kwargs)
+    return train_loader, val_loader
+
+
+def get_test_loader_mnist(
+    data_dir: str, batch_size: int
+) -> torch.utils.data.DataLoader:
+    """Test loader of the MNIST dataset.
+
+    Args:
+        data_dir (str): Directory of the MNIST dataset.
+        batch_size (int): Size of the batches of the test loader.
+
+    Returns:
+        torch.utils.data.DataLoader: Test loader.
+    """
+    dataset = preprocess_mnist(data_dir, train=False)
+    loader = DataLoader(dataset, shuffle=False, batch_size=batch_size)
+    return loader
+
+
+def read_digit_image(image_file: FILE_LIKE) -> torch.FloatTensor:
+    """Loads a single digit image from a file.
+
+    Center crops and resizes the image to 28 by 28 pixels.  Also inverts the image if
+    there are more bright than dark pixels.
+
+    Args:
+        image_file (FILE_LIKE): The image file.
+
+    Returns:
+        torch.FloatTensor: The preprocessed digit image.
+    """
+    image = read_image(image_file, mode=ImageReadMode.GRAY)
+    transform = transforms.Compose(
+        [
+            transforms.CenterCrop(min(image.size()[1:])),
+            transforms.Resize([28, 28]),
+        ]
+    )
+    # Standardize image
+    image = transform(image).type(torch.FloatTensor)
+    image = transforms.Normalize(image.mean(), image.std())(image)
+    # Check if we need to invert the image
+    if (image > 0).count_nonzero() > image.numel() / 2:
+        # More bright than dark pixels
+        image = transforms.functional.invert(image)
+    return image