0-utils-data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e6bd4799",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |default_exp utils.data\n",
    "# |export\n",
    "from typing import Tuple\n",
    "import torch\n",
    "from torch.utils import data\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3863535",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bedb1b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |export\n",
    "def load_iris(*, force_single_precision=False) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Load the Iris dataset and return the data and targets.\n",
    "\n",
    "    Args:\n",
    "        force_single_precision: Whether to force the data to be single precision.\n",
    "    Returns:\n",
    "        Tuple[torch.Tensor, torch.Tensor]: The data and targets.\n",
    "    \"\"\"\n",
    "    from sklearn import datasets\n",
    "\n",
    "    iris = datasets.load_iris()\n",
    "    data = torch.from_numpy(iris[\"data\"])\n",
    "    if force_single_precision:\n",
    "        data = data.float()\n",
    "\n",
    "    targets = torch.from_numpy(iris[\"target\"]).to(torch.long)\n",
    "    return data, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c1a90b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |export\n",
    "def get_mnist_datasets(cache_path: str) -> Tuple[data.Dataset, data.Dataset]:\n",
    "    \"\"\"\n",
    "    Get the MNIST datasets.\n",
    "\n",
    "    Args:\n",
    "        cache_path: The path to the cache directory.\n",
    "    Returns:\n",
    "        Tuple[data.Dataset, data.Dataset]: The MNIST train and test datasets.\n",
    "    \"\"\"\n",
    "\n",
    "    transform = transforms.Compose(\n",
    "        [\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(\n",
    "                (0.1307,), (0.3081,)\n",
    "            ),  # normalization, see https://discuss.pytorch.org/t/normalization-in-the-mnist-example/457\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    mnist_train_set = datasets.MNIST(\n",
    "        root=cache_path, train=True, download=True, transform=transform\n",
    "    )\n",
    "    mnist_test_set = datasets.MNIST(\n",
    "        root=cache_path, train=False, download=True, transform=transform\n",
    "    )\n",
    "    return mnist_train_set, mnist_test_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f7a0e9e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# |export\n",
    "def get_fashion_mnist_datasets(cache_path: str) -> Tuple[data.Dataset, data.Dataset]:\n",
    "    \"\"\"\n",
    "    Get the Fashion MNIST datasets.\n",
    "\n",
    "    Args:\n",
    "        cache_path: The path to the cache directory.\n",
    "    Returns:\n",
    "        Tuple[data.Dataset, data.Dataset]: The Fashion MNIST train and test datasets.\n",
    "    \"\"\"\n",
    "\n",
    "    fmnist_train_set = datasets.FashionMNIST(\n",
    "        root=cache_path, train=True, download=True, transform=transforms.ToTensor()\n",
    "    )\n",
    "    fmnist_test_set = datasets.FashionMNIST(\n",
    "        root=cache_path, train=False, download=True, transform=transforms.ToTensor()\n",
    "    )\n",
    "    return fmnist_train_set, fmnist_test_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4aebdccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "cwd = os.getcwd()\n",
    "cache_path = os.path.join(cwd, \"datasets\")\n",
    "\n",
    "fmnist_trainset, fmnist_testset = get_fashion_mnist_datasets(cache_path)\n",
    "mnist_trainset, mnist_testset = get_mnist_datasets(cache_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}