forward_forward.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a548214",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "from torch.optim import Adam\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision.transforms import Compose, ToTensor, Normalize, Lambda\n",
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "\n",
    "DEVICE = 'cuda'\n",
    "# DEVICE = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fceae2ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def MNIST_loaders(train_batch_size=20000, test_batch_size=10000):\n",
    "    transform = Compose([\n",
    "        ToTensor(),\n",
    "        Normalize((0.1307,), (0.3081,)),\n",
    "        Lambda(lambda x: torch.flatten(x))])\n",
    "\n",
    "    train_loader = DataLoader(\n",
    "        MNIST('./data/', train=True,\n",
    "              download=True,\n",
    "              transform=transform),\n",
    "        batch_size=train_batch_size, shuffle=True)\n",
    "\n",
    "    test_loader = DataLoader(\n",
    "        MNIST('./data/', train=False,\n",
    "              download=True,\n",
    "              transform=transform),\n",
    "        batch_size=test_batch_size, shuffle=False)\n",
    "\n",
    "    return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bce8f24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def overlay_y_on_x(x, y):\n",
    "    x_ = x.clone()\n",
    "    x_[:, :10] *= 0.0\n",
    "    x_[range(x.shape[0]), y] = x.max()\n",
    "    return x_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "356d9743",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, dims):\n",
    "        super().__init__()\n",
    "        self.layers = []\n",
    "        for d in range(len(dims) - 1):\n",
    "            self.layers += [Layer(dims[d], dims[d + 1]).to(DEVICE)]\n",
    "\n",
    "    def predict(self, x):\n",
    "        goodness_per_label = []\n",
    "        for label in range(10):\n",
    "            h = overlay_y_on_x(x, label)\n",
    "            goodness = []\n",
    "            \n",
    "            for layer in self.layers:\n",
    "                h = layer(h)\n",
    "                goodness += [h.pow(2).mean(1)]\n",
    "            goodness_per_label += [sum(goodness).unsqueeze(1)]\n",
    "        goodness_per_label = torch.cat(goodness_per_label, 1)\n",
    "        return goodness_per_label.argmax(1)\n",
    "\n",
    "    def train(self, x_pos, x_neg):\n",
    "        h_pos, h_neg = x_pos, x_neg\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            h_pos, h_neg = layer.train(h_pos, h_neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8241944",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Layer(nn.Linear):\n",
    "    def __init__(self, in_features, out_features,\n",
    "                 bias=True, device=None, dtype=None):\n",
    "        super().__init__(in_features, out_features, bias, device, dtype)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.opt = Adam(self.parameters(), lr=0.03)\n",
    "        self.threshold = 3.0\n",
    "        self.num_epochs = 1000\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_direction = x / (x.pow(2).sum(dim=1).sqrt().reshape((x.shape[0], 1)) + 1e-4)\n",
    "        \n",
    "        return self.relu(\n",
    "            torch.mm(x_direction, self.weight.T) +\n",
    "            self.bias.unsqueeze(0))\n",
    "\n",
    "    def train(self, x_pos, x_neg):\n",
    "        for i in range(self.num_epochs):\n",
    "            g_pos = self.forward(x_pos).pow(2).mean(1)\n",
    "            g_neg = self.forward(x_neg).pow(2).mean(1)\n",
    "            \n",
    "            pos_loss = -g_pos + self.threshold\n",
    "            neg_loss = g_neg - self.threshold\n",
    "            \n",
    "            loss = torch.log(1 + torch.exp(torch.cat([\n",
    "                                    pos_loss,\n",
    "                                    neg_loss]))).mean()\n",
    "            self.opt.zero_grad()\n",
    "\n",
    "            loss.backward()\n",
    "            self.opt.step()\n",
    "\n",
    "        return self.forward(x_pos).detach(), self.forward(x_neg).detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0149ec51",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9363c8fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_negative_data(x, y):\n",
    "    y_neg = y.clone()\n",
    "    for idx, y_samp in enumerate(y):\n",
    "        allowed_indices = [i for i in range(10)]\n",
    "        allowed_indices.pop(y_samp.item())\n",
    "        y_neg[idx] = torch.tensor(np.random.choice(allowed_indices)).to(DEVICE)\n",
    "    \n",
    "    return overlay_y_on_x(x, y_neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e2531ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1234)\n",
    "train_loader, test_loader = MNIST_loaders()\n",
    "\n",
    "net = Net([784, 512, 512])\n",
    "\n",
    "for x, y in tqdm(train_loader):\n",
    "    x, y = x.to(DEVICE), y.to(DEVICE)\n",
    "\n",
    "    x_pos = overlay_y_on_x(x, y)\n",
    "    x_neg = generate_negative_data(x, y)\n",
    "    \n",
    "    net.train(x_pos, x_neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cff9cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.save(net, 'net.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57600750",
   "metadata": {},
   "outputs": [],
   "source": [
    "# net = torch.load('net.pth')\n",
    "# train_loader, test_loader = MNIST_loaders()\n",
    "# torch.manual_seed(1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd6a1038",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_err = round((1.0 - net.predict(x).eq(y).float().mean().item()) * 100, 2)\n",
    "print('train error:', train_err)\n",
    "\n",
    "x_te, y_te = next(iter(test_loader))\n",
    "x_te, y_te = x_te.to(DEVICE), y_te.to(DEVICE)\n",
    "\n",
    "test_err = round((1.0 - net.predict(x_te).eq(y_te).float().mean().item()) * 100, 2)\n",
    "print('test error:', test_err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93717577",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9eaf985",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(y_te[:1])\n",
    "img = x_te[:1,:].reshape((28,28)).cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4fc4a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d41a5c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "net.layers[0].weight.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b57dc0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(pos_img.reshape((28,28)).cpu())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "962ce841",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_img = overlay_y_on_x(x_te[:1,:], 7)\n",
    "neg_img = overlay_y_on_x(x_te[:1,:], 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4abf075",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_activations(img, title):\n",
    "    activations = []\n",
    "\n",
    "    activ = img\n",
    "\n",
    "    for layer in net.layers:\n",
    "        activ = layer(activ)\n",
    "        activations.append(activ)\n",
    "\n",
    "    fig = plt.figure(figsize=(6,4))\n",
    "    fig.suptitle(title)\n",
    "    columns = 2\n",
    "    \n",
    "    for i, image in enumerate(activations):\n",
    "        plt.subplot(1, columns, i+1)\n",
    "        plt.imshow(image.detach().cpu().reshape((32,-1)))\n",
    "        \n",
    "    return activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "510e6201",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_act = show_activations(pos_img, \"Положительные данные\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ba7139",
   "metadata": {},
   "outputs": [],
   "source": [
    "neg_act = show_activations(neg_img, \"Негативные данные\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53683ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "neg_act[0].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "874cf602",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffs = []\n",
    "\n",
    "for i, act in enumerate(pos_act):\n",
    "    diff = act - neg_act[i]\n",
    "    diffs.append(diff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25c68004",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,4))\n",
    "columns = 3\n",
    "for i, image in enumerate(diffs):\n",
    "    plt.subplot(1, columns, i+1)\n",
    "#     plt.imshow(image.detach().cpu().reshape((32,-1)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}