{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "hourly-carrier",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10000, 1, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "# Load data (MNIST)\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "batch_size_train = 100000\n",
    "batch_size_test = 20000\n",
    "log_interval = 10\n",
    "\n",
    "\n",
    "# num_samples = 128 # number of training samples\n",
    "num_samples = 1024 * 8\n",
    "num_samples_test = 100 # number of test samples\n",
    "\n",
    "new_dim1 = 28 * 1 # first dimension\n",
    "new_dim2 = 28 * 1 # second dimension\n",
    "\n",
    "\n",
    "old_dim = 28 # MNIST original dimension\n",
    "\n",
    "\n",
    "random_seed = 1\n",
    "torch.backends.cudnn.enabled = False\n",
    "torch.manual_seed(random_seed)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root = './', train=True, download=True,\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor(),\n",
    "                               torchvision.transforms.Normalize(\n",
    "                                 (0.1307,), (0.3081,))\n",
    "                             ])),\n",
    "    batch_size=batch_size_train, shuffle=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root = './', train=False, download=True,\n",
    "                             transform=torchvision.transforms.Compose([\n",
    "                               torchvision.transforms.ToTensor(),\n",
    "                               torchvision.transforms.Normalize(\n",
    "                                 (0.1307,), (0.3081,))\n",
    "                             ])),\n",
    "    batch_size=batch_size_test, shuffle=True)\n",
    "\n",
    "examples = enumerate(test_loader)\n",
    "batch_idx, (example_data, example_targets) = next(examples)\n",
    "\n",
    "print(example_data.shape)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "thorough-array",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "# Grayscale version of CIFAR-10\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "batch_size_train = 10000\n",
    "batch_size_test = 20000\n",
    "log_interval = 10\n",
    "\n",
    "# num_samples = 1024 * 4 # number of training samples\n",
    "num_samples = 512 # number of training samples\n",
    "\n",
    "\n",
    "\n",
    "# num_samples = 128 # number of training samples\n",
    "num_samples_test = 100 # number of test samples\n",
    "\n",
    "new_dim1 = 32 * 1 # first dimension\n",
    "new_dim2 = 32 * 1 # second dimension\n",
    "\n",
    "\n",
    "old_dim = 32 # MNIST original dimension\n",
    "\n",
    "\n",
    "transform = transforms.Compose(\n",
    "    [transforms.Grayscale(num_output_channels=1),\n",
    "                                    transforms.ToTensor(), \n",
    "                                    transforms.Normalize((0.5, ), (0.5, ))])\n",
    "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size_train,\n",
    "                                          shuffle=True, num_workers=1)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(root='./data', train=False,\n",
    "                                       download=True, transform=transform)\n",
    "test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size_test,\n",
    "                                         shuffle=False, num_workers=1)\n",
    "\n",
    "# classes = ('plane', 'car', 'bird', 'cat',\n",
    "#            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "excellent-coupon",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10000, 1, 32, 32])\n",
      "torch.Size([10000])\n"
     ]
    }
   ],
   "source": [
    "examples = enumerate(test_loader)\n",
    "batch_idx, (example_data, example_targets) = next(examples)\n",
    "\n",
    "print(example_data.shape)\n",
    "print(example_targets.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "insured-cincinnati",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded\n",
      "Created training data\n",
      "Created test data\n"
     ]
    }
   ],
   "source": [
    "# Create training data and test data for binary classification task\n",
    "\n",
    "from tqdm import tqdm \n",
    "training_data = enumerate(train_loader)\n",
    "test_data = enumerate(test_loader)\n",
    "(batch_id1, (data_tr_old, target_tr_old)) = next(training_data)\n",
    "(batch_id2, (data_test_old, target_test_old)) = next(test_data)\n",
    "print(\"loaded\")\n",
    "\n",
    "\n",
    "# Classes\n",
    "# class1 = [0, 1, 2, 3, 4]\n",
    "# class2 = [5, 6, 7, 8, 9]\n",
    "class1 = [0]\n",
    "class2 = [2]\n",
    "# Pad the image on the right and on the bottom\n",
    "def filter_digits(data_old, target_old, num, threshold):\n",
    "    data = torch.zeros([num, 1, new_dim1, new_dim2])\n",
    "    target = torch.zeros([num])\n",
    "\n",
    "    num_samples_per_class = int(num / 2)\n",
    "    sample1 = 0\n",
    "    sample2 = 0\n",
    "    attempt = 0\n",
    "    sample = 0\n",
    "\n",
    "    while (sample1 < num_samples_per_class) or (sample2 < num_samples_per_class):\n",
    "        if attempt == threshold:\n",
    "            print(\"FAILED: need more samples in batch\")\n",
    "            return (data, target)\n",
    "        # Balance classes\n",
    "        is_sample_for_class1 = (target_old[attempt] in class1) and (sample1 < num_samples_per_class)\n",
    "        is_sample_for_class2 = (target_old[attempt] in class2) and (sample2 < num_samples_per_class)\n",
    "        if is_sample_for_class1 or is_sample_for_class2: \n",
    "            avg = torch.mean(data_old[attempt][0])\n",
    "            target[sample] = target_old[attempt]\n",
    "            for i in range(old_dim):\n",
    "                for j in range(old_dim):\n",
    "                    data[sample][0][i][j] = data_old[attempt][0][i][j] - avg\n",
    "            if is_sample_for_class1:\n",
    "                sample1 += 1\n",
    "            if is_sample_for_class2:\n",
    "                sample2 += 1\n",
    "          # Augment sample counts\n",
    "            sample += 1\n",
    "        attempt += 1\n",
    "\n",
    "    target.apply_(lambda x: 1 if (x in class1) else -1)\n",
    "\n",
    "    data = data.float()\n",
    "    target = target.float()\n",
    "    return (data, target)\n",
    "\n",
    "\n",
    "# Training data and test data\n",
    "(data_tr, target_tr) = filter_digits(data_tr_old, target_tr_old, num_samples, batch_size_train)\n",
    "print(\"Created training data\")\n",
    "(data_test, target_test) = filter_digits(data_test_old, target_test_old, num_samples_test, batch_size_test)\n",
    "print(\"Created test data\")\n",
    "\n",
    "\n",
    "torch.save(data_tr, \"training-test-data/training_data_nonlinear.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_tr, \"training-test-data/training_targets_nonlinear.txt\")\n",
    "\n",
    "\n",
    "torch.save(data_test, \"training-test-data/test_data_nonlinear.txt\")\n",
    "\n",
    "\n",
    "torch.save(target_test, \"training-test-data/test_targets_nonlinear.txt\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "discrete-sensitivity",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5\n"
     ]
    }
   ],
   "source": [
    "# Check that classes are balanced\n",
    "counter = 0\n",
    "for i in range(num_samples):\n",
    "    if target_tr[i] == 1:\n",
    "        counter += 1\n",
    "print(counter/num_samples)\n",
    "\n",
    "counter = 0\n",
    "for i in range(num_samples_test):\n",
    "    if target_test[i] == 1:\n",
    "        counter += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "peripheral-passion",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2, 3, 4]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "instructional-redhead",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
