SIAM CSE23 minitutorial PyG Part1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Hands-on Introduction to Geometric Deep Learning, with Examples in PyTorch Geometric - I"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SIAM-CSE, Amsterdam\n",
    "### February 26 - March 3, 2023"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Outline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Datasets\n",
    "- Models\n",
    "- Losses\n",
    "- Optimizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Official resources:\n",
    "* [Deep Learning with PyTorch: a 60 Minute Blitz](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html)\n",
    "* [PyTorch documentation](https://pytorch.org/docs/stable/index.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here no details, mostly visualization, with linear regression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "plt.rcParams.update({'font.size': 16})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand(20, 5)\n",
    "x\n",
    "((x.norm(dim=1) - 1).abs() < 1e-10).float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A raw dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data are obtained from the model \n",
    "$$\n",
    "x \\mapsto y:=Ax + b\n",
    "$$\n",
    "with $x\\in\\mathbb R^{input\\_dim}$ and $y\\in\\mathbb R^{output\\_dim}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dim = 1\n",
    "output_dim = 1\n",
    "\n",
    "A = 2 * np.random.rand(output_dim, input_dim) - 1\n",
    "b = 2 * np.random.rand(output_dim) - 1\n",
    "\n",
    "true_model = lambda x: A @ x + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get a training set by randomly sampling and noisy observations\n",
    "\\begin{align*}\n",
    "&x_i \\sim U([-1, 1])\\\\\n",
    "&y_i = A x_i + b + \\nu_i\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train = 1000\n",
    "noise_level = 0.04\n",
    "\n",
    "# Generate a random set of n_train samples\n",
    "X_train = np.random.rand(n_train, input_dim)\n",
    "y_train = np.array([true_model(x) for x in X_train])\n",
    "\n",
    "# Add some noise\n",
    "y_train += noise_level * np.random.standard_normal(size=y_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    fig = plt.figure()\n",
    "    fig.clf()\n",
    "    ax = fig.gca()\n",
    "    ax.plot(X_train, y_train, '.')\n",
    "    ax.grid(True)\n",
    "    ax.set_xlabel('X_train')\n",
    "    ax.set_ylabel('y_train')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch `Dataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% Dataset to manage vector to vector data\n",
    "class VectorialDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, input_data, output_data):\n",
    "        super(VectorialDataset, self).__init__()\n",
    "        self.input_data = torch.tensor(input_data.astype('f'))\n",
    "        self.output_data = torch.tensor(output_data.astype('f'))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.input_data.shape[0]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "        sample = (self.input_data[idx, :], \n",
    "                  self.output_data[idx, :])  \n",
    "        return sample "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_set = VectorialDataset(input_data=X_train, output_data=y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(training_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_set[10:12]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here almost useless, but think about e.g. images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch `DataLoader`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 120\n",
    "train_loader = torch.utils.data.DataLoader(training_set, \n",
    "                                           batch_size=batch_size, \n",
    "                                           shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The role of `batch_size`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, batch in enumerate(train_loader):\n",
    "    print('Batch n. %2d: input size=%s, output size=%s' % (idx+1, batch[0].shape, batch[1].shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "8 * 120"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The role of `shuffle`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_batch = []\n",
    "\n",
    "for epoch in range(2):\n",
    "    for idx, batch in enumerate(train_loader):\n",
    "        if idx == 0:\n",
    "            first_batch.append(batch)\n",
    "        \n",
    "np.c_[X_train[:batch_size], first_batch[0][0].numpy(), first_batch[1][0].numpy()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We implement a linear model $$x \\mapsto model(x):=A x + b$$ \n",
    "\n",
    "with $A\\in \\mathbb{R}^{input\\_dim\\times output\\_dim}$, $b\\in\\mathbb{R}^{output\\_dim}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "\n",
    "#%% Linear layer\n",
    "class LinearModel(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim):\n",
    "        super(LinearModel, self).__init__()\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "\n",
    "        self.linear = nn.Linear(self.input_dim, self.output_dim, bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.linear(x)\n",
    "        return out\n",
    "    \n",
    "    def reset(self):\n",
    "        self.linear.reset_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LinearModel(input_dim, output_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See [pytorch.org/docs/stable/nn.html](https://pytorch.org/docs/stable/nn.html) for many other layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.linear.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.linear.bias"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The `forward` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(5, input_dim)\n",
    "model.forward(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[model.linear.weight @ xx + model.linear.bias for xx in x]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    fig = plt.figure()\n",
    "    fig.clf()\n",
    "    ax = fig.gca()\n",
    "    ax.plot(training_set.input_data, training_set.output_data, '.')\n",
    "    ax.plot(training_set.input_data, model.forward(training_set.input_data).detach().numpy(), '.')\n",
    "    ax.grid(True)\n",
    "    ax.set_xlabel('X_train')\n",
    "    ax.legend(['y_train', 'model(X_train)'])    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The MSE loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MSE loss is\n",
    "$$\n",
    "L(y, y') = \\|y-y'\\|_2^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "loss_fun = nn.MSELoss(reduction='mean')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More at [pytorch.org/docs/stable/nn.html#loss-functions](https://pytorch.org/docs/stable/nn.html#loss-functions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor(np.array([1, 2, 1]).astype('f'))\n",
    "z = torch.tensor(np.array([0, 0, 0]).astype('f'))\n",
    "loss_fun(x, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Giving a score to the model (parameters) given the training set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss for one data pair $(x_i, y_i)$ is \n",
    "$$\n",
    "L(model(x_i), y_i)\n",
    "$$\n",
    "\n",
    "The cumulative (average) loss on the training set is \n",
    "$$\n",
    "L(X_{train}, y_{train}):=\\frac{1}{n_{train}} \\sum_{i=1}^{n_{train}} L(model(x_i), y_i)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "\n",
    "    state_dict = model.state_dict()\n",
    "\n",
    "    ww, bb = np.meshgrid(np.linspace(-2, 2, 30), np.linspace(-2, 2, 30))\n",
    "\n",
    "    loss_values = 0 * ww\n",
    "    for i in range(ww.shape[0]):\n",
    "        for j in range(ww.shape[1]):\n",
    "            state_dict['linear.weight'] = torch.tensor([[ww[i, j]]])\n",
    "            state_dict['linear.bias'] = torch.tensor([bb[i, j]])\n",
    "            model.load_state_dict(state_dict)\n",
    "            loss_values[i, j] = loss_fun(model.forward(training_set.input_data),  training_set.output_data)\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 8))\n",
    "    fig.clf()\n",
    "    ax = fig.gca()\n",
    "    levels = np.logspace(np.log(np.min(loss_values)), np.log(np.max(loss_values)), 20)\n",
    "    c=ax.contourf(ww, bb, loss_values, levels=levels, norm=colors.LogNorm())\n",
    "    plt.colorbar(c)\n",
    "    ax.plot(A[0], b, 'r*', markersize=10)\n",
    "    ax.set_ylabel('bias')\n",
    "    ax.set_xlabel('weight')\n",
    "    \n",
    "    ax.grid(True)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Goal: minimize the loss given the data:\n",
    "$$\n",
    "\\min\\limits_{par \\in model} L(X_{train}, y_{train})\n",
    "=\\min\\limits_{par \\in model} \\frac{1}{n_{train}} \\sum_{i=1}^{n_{train}} L(x_i, model(x_i))\n",
    "$$\n",
    "\n",
    "In this case $par =\\{A, b\\}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterative gradient based optimization:\n",
    "\\begin{align*}\n",
    "par^{(0)} &= par_0\\\\\n",
    "par^{(k+1)} &=  par^{(k)} - \\eta^{(k)} \\nabla_{par} L(X_{train}, y_{train})\n",
    "\\end{align*}    \n",
    "with learning rate $\\eta^{(k)}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model is \n",
    "$$\n",
    "model(x) = A x + b\n",
    "$$\n",
    "\n",
    "With $input\\_dim = output\\_dim = 1$ we have $A:=a\\in\\mathbb{R},\\ b\\in\\mathbb{R}$. \n",
    "\n",
    "For a data pair $(x, y)$ the loss is\n",
    "\\begin{align*}\n",
    "L(x, y) \n",
    "&= \\|model(x)-y\\|_2^2\\\\ \n",
    "&= (model(x)-y)^2\\\\\n",
    "&= (a x + b - y)^2\\\\\n",
    "&= a^2 x^2 + 2 a b x - 2 a xy + b^2 -2 b y + y^2.\n",
    "\\end{align*}\n",
    "\n",
    "We can compute\n",
    "\\begin{align*}\n",
    "\\nabla_a L(x, y) \n",
    "&= 2 a x^2 + 2 b x - 2 xy\n",
    "=2 x (a x + b - y)\\\\\n",
    "\\nabla_b L(x, y)\n",
    "&= 2 a x + 2 b - 2 y\n",
    "= 2 (ax + b - y).\n",
    "\\end{align*}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1, input_dim)\n",
    "y =  torch.randn(1, output_dim)\n",
    "\n",
    "model.zero_grad()\n",
    "loss = loss_fun(model.forward(x),  y)\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    print(model.linear.weight.grad)\n",
    "    print(2 * x * (model.linear.weight * x + model.linear.bias - y))\n",
    "    \n",
    "    print(model.linear.bias.grad)\n",
    "    print(2 * (model.linear.weight * x + model.linear.bias - y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Handmade optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "\n",
    "    num_iter = 200\n",
    "    lr = 0.5 # 0.01\n",
    "\n",
    "    train_hist = {}\n",
    "    train_hist['weight'] = []\n",
    "    train_hist['bias'] = []\n",
    "\n",
    "    model.reset()\n",
    "    state_dict = model.state_dict()\n",
    "\n",
    "    for _ in range(num_iter):\n",
    "\n",
    "        model.zero_grad()\n",
    "        loss = loss_fun(model.forward(training_set.input_data), training_set.output_data)\n",
    "        loss.backward()\n",
    "\n",
    "        w = model.linear.weight.item()\n",
    "        b = model.linear.bias.item()\n",
    "\n",
    "        dw = model.linear.weight.grad.item()\n",
    "        db = model.linear.bias.grad.item()\n",
    "\n",
    "        state_dict['linear.weight'] += torch.tensor([-lr * dw])\n",
    "        state_dict['linear.bias'] += torch.tensor([-lr * db])\n",
    "        model.load_state_dict(state_dict)\n",
    "\n",
    "        train_hist['weight'].append(w)\n",
    "        train_hist['bias'].append(b)\n",
    "\n",
    "    for label in train_hist:\n",
    "        train_hist[label] = np.array(train_hist[label])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    fig = plt.figure(figsize=(8, 8))\n",
    "    fig.clf()\n",
    "    ax = fig.gca()\n",
    "    levels = np.logspace(np.log(np.min(loss_values)), np.log(np.max(loss_values)), 20)\n",
    "    ax.contourf(ww, bb, loss_values, levels=levels, norm=colors.LogNorm())\n",
    "    ax.set_xlabel('weight')\n",
    "    ax.set_ylabel('bias')\n",
    "    ax.grid(True)\n",
    "    ax.set_xlim(-2, 2) \n",
    "    ax.set_ylim(-2, 2) \n",
    "    \n",
    "    ax.plot(train_hist['weight'], train_hist['bias'], '.-b')\n",
    "    ax.plot(A[0], b, 'r*', markersize=10)\n",
    "\n",
    "    #ax.legend(['optim', '(A, b)'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stochastic GD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Avoid loading the full training set\n",
    "* Avoid evaluating the model on the full training set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At every step, compute the loss on a batch of data $(X^{(k)}, y^{(k)})\\sim (X_{train}, y_{train})$: \n",
    "\n",
    "\\begin{align*}\n",
    "L(X^{(k)}, y^{(k)}) := \\sum_{(x, y) \\in (X^{(k)}, y^{(k)})} L(x, model(x)).\n",
    "\\end{align*}    \n",
    "\n",
    "and then update with this approximated gradient:\n",
    "\\begin{align*}\n",
    "par^{(k+1)} &=  par^{(k)} - \\eta^{(k)} \\nabla_{par} L(X^{(k)}, y^{(k)}). \n",
    "\\end{align*}    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch `optim`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.1\n",
    "weight_decay = 5e-4\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More at [pytorch.org/docs/stable/optim.html](https://pytorch.org/docs/stable/optim.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_hist = {}\n",
    "train_hist['loss'] = []\n",
    "\n",
    "if input_dim == output_dim == 1:\n",
    "    train_hist['weight'] = []\n",
    "    train_hist['bias'] = []\n",
    "\n",
    "# Initialize training\n",
    "model.reset()\n",
    "model.train()\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    for idx, batch in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        loss = loss_fun(model.forward(batch[0]),  batch[1])\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_hist['loss'].append(loss.item())\n",
    "        if input_dim == output_dim == 1:\n",
    "            train_hist['weight'].append(model.linear.weight.item())\n",
    "            train_hist['bias'].append(model.linear.bias.item())\n",
    "        \n",
    "        print('[Epoch %4d/%4d] [Batch %4d/%4d] Loss: % 2.2e' % (epoch + 1, n_epochs, \n",
    "                                                                idx + 1, len(train_loader), \n",
    "                                                                loss.item()))\n",
    "        \n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    n_test = 500\n",
    "    X_test = np.random.rand(n_test, input_dim)\n",
    "    y_pred = []\n",
    "\n",
    "    state_dict = model.state_dict()\n",
    "\n",
    "    for idx in range(len(train_hist['weight'])):\n",
    "        state_dict['linear.weight'] = torch.tensor([[train_hist['weight'][idx]]])\n",
    "        state_dict['linear.bias'] = torch.tensor([train_hist['bias'][idx]])\n",
    "        model.load_state_dict(state_dict)\n",
    "\n",
    "        y_pred.append(model.forward(torch.tensor(X_test.astype('f'))).detach().numpy())    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_dim == output_dim == 1:\n",
    "    fig = plt.figure(figsize=(15, 5))\n",
    "    fig.clf()\n",
    "    \n",
    "    ax = fig.add_subplot(1, 3, 1)\n",
    "    levels = np.logspace(np.log(np.min(loss_values)), np.log(np.max(loss_values)), 20)\n",
    "    ax.contourf(ww, bb, loss_values, levels=levels, norm=colors.LogNorm())\n",
    "    ax.plot(train_hist['weight'], train_hist['bias'], '.-b')\n",
    "    ax.plot(A[0], b, 'r*', markersize=10)\n",
    "    ax.set_xlabel('weight')\n",
    "    ax.set_ylabel('bias')\n",
    "    #ax.legend(['optim', '(A, b)'])\n",
    "    ax.grid(True)\n",
    "    ax.set_xlim(-2, 2) \n",
    "    ax.set_ylim(-2, 2) \n",
    "    \n",
    "    ax = fig.add_subplot(1, 3, 2)\n",
    "    ax.loglog(np.abs(train_hist['loss']))\n",
    "    ax.set_xlabel('Iter')\n",
    "    ax.set_ylabel('Loss')\n",
    "    ax.grid(True)\n",
    "    \n",
    "    ax = fig.add_subplot(1, 3, 3)\n",
    "    ax.plot(X_train, y_train, '.')\n",
    "    a=ax.plot(X_test, y_pred[0], '-', alpha=0.1)\n",
    "    for y in y_pred[1:]:\n",
    "        ax.plot(X_test, y, '-', alpha=0.1, color=a[0].get_color())\n",
    "    ax.plot(X_test, y_pred[-1], 'k')\n",
    "    ax.grid(True)   \n",
    "\n",
    "    fig.tight_layout()\n",
    "\n",
    "else:\n",
    "    fig = plt.figure()\n",
    "    fig.clf()\n",
    "    ax = fig.gca()\n",
    "    ax.loglog(np.abs(train_hist['loss']))\n",
    "    ax.set_xlabel('Iter')\n",
    "    ax.set_ylabel('Loss')\n",
    "    ax.grid(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 4
}