4-training-on-Darcy-Flow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "012a357f-8533-482c-823d-a4587c49e726",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import wandb\n",
    "import sys\n",
    "from configmypy import ConfigPipeline, YamlConfig, ArgparseConfig\n",
    "from neuralop import get_model\n",
    "from neuralop import Trainer\n",
    "from neuralop.training import setup\n",
    "from neuralop.datasets import load_darcy_pt\n",
    "from neuralop.utils import get_wandb_api_key, count_params\n",
    "from neuralop import LpLoss, H1Loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b6358b5-78d1-4baf-8928-6bb49b150680",
   "metadata": {},
   "source": [
    "# Loading the configuration\n",
    "\n",
    "You can open the yaml file in config/darcy_config in the same folder as this notebook to inspect the parameters and change them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4503f065-4063-4a4f-b00f-06a7c3a88e27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the configuration\n",
    "config_name = 'default'\n",
    "pipe = ConfigPipeline([YamlConfig('./darcy_config.yaml', config_name='default', config_folder='./config'),\n",
    "                      ])\n",
    "config = pipe.read_conf()\n",
    "config_name = pipe.steps[-1].config_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e95d820d-9578-4ad7-80b4-05a5771f1642",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Here we just setup pytorch and print the configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "46066d9f-21a3-4aab-b6e1-f7f38e05f88b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set-up distributed communication, if using\n",
    "device, is_logger = setup(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "26d599f9-6463-4056-9a4d-72c01d05298e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###############################\n",
      "#####    CONFIGURATION    #####\n",
      "###############################\n",
      "\n",
      "Steps:\n",
      "------\n",
      " (1) YamlConfig with config_file=./darcy_config.yaml, config_name=default, config_folder=./config\n",
      "\n",
      "-------------------------------\n",
      "\n",
      "Configuration:\n",
      "--------------\n",
      "\n",
      "n_params_baseline=None\n",
      "verbose=True\n",
      "arch=tfno2d\n",
      "distributed.use_distributed=False\n",
      "tfno2d.data_channels=3\n",
      "tfno2d.n_modes_height=32\n",
      "tfno2d.n_modes_width=32\n",
      "tfno2d.hidden_channels=64\n",
      "tfno2d.projection_channels=256\n",
      "tfno2d.n_layers=4\n",
      "tfno2d.domain_padding=None\n",
      "tfno2d.domain_padding_mode=one-sided\n",
      "tfno2d.fft_norm=forward\n",
      "tfno2d.norm=group_norm\n",
      "tfno2d.skip=linear\n",
      "tfno2d.implementation=factorized\n",
      "tfno2d.separable=0\n",
      "tfno2d.preactivation=0\n",
      "tfno2d.use_mlp=1\n",
      "tfno2d.mlp.expansion=0.5\n",
      "tfno2d.mlp.dropout=0\n",
      "tfno2d.factorization=None\n",
      "tfno2d.rank=1.0\n",
      "tfno2d.fixed_rank_modes=None\n",
      "tfno2d.dropout=0.0\n",
      "tfno2d.tensor_lasso_penalty=0.0\n",
      "tfno2d.joint_factorization=False\n",
      "opt.n_epochs=150\n",
      "opt.learning_rate=0.005\n",
      "opt.training_loss=h1\n",
      "opt.weight_decay=0.0001\n",
      "opt.amp_autocast=False\n",
      "opt.scheduler_T_max=300\n",
      "opt.scheduler_patience=5\n",
      "opt.scheduler=CosineAnnealingLR\n",
      "opt.step_size=50\n",
      "opt.gamma=0.5\n",
      "data.folder=/data/darcy_flow/\n",
      "data.batch_size=32\n",
      "data.n_train=3000\n",
      "data.train_resolution=32\n",
      "data.n_tests=[500, 500]\n",
      "data.test_resolutions=[32, 64]\n",
      "data.test_batch_sizes=[32, 32]\n",
      "data.positional_encoding=True\n",
      "data.encode_input=True\n",
      "data.encode_output=False\n",
      "patching.levels=0\n",
      "patching.padding=0\n",
      "patching.stitching=False\n",
      "wandb.log=False\n",
      "wandb.log_test_interval=1\n",
      "\n",
      "###############################\n"
     ]
    }
   ],
   "source": [
    "# Make sure we only print information when needed\n",
    "config.verbose = config.verbose and is_logger\n",
    "\n",
    "#Print config to screen\n",
    "if config.verbose and is_logger:\n",
    "    pipe.log()\n",
    "    sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1339c794-3e1c-469b-b0a0-cf968fc1dfa1",
   "metadata": {},
   "source": [
    "# Loading the data \n",
    "\n",
    "We train in one resolution and test in several resolutions to show the zero-shot super-resolution capabilities of neural-operators. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3515a85a-40fc-4223-9cdb-8768de37d6e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UnitGaussianNormalizer init on 3000, reducing over [0, 1, 2, 3], samples of shape [1, 32, 32].\n",
      "   Mean and std of shape torch.Size([1, 1, 1]), eps=1e-05\n",
      "Loading test db at resolution 64 with 500 samples and batch-size=32\n"
     ]
    }
   ],
   "source": [
    "# Loading the Darcy flow training set in 32x32 resolution, test set in 32x32 and 64x64 resolutions\n",
    "train_loader, test_loaders, output_encoder = load_darcy_pt(\n",
    "        config.data.folder, train_resolution=config.data.train_resolution, n_train=config.data.n_train, batch_size=config.data.batch_size, \n",
    "        positional_encoding=config.data.positional_encoding,\n",
    "        test_resolutions=config.data.test_resolutions, n_tests=config.data.n_tests, test_batch_sizes=config.data.test_batch_sizes,\n",
    "        encode_input=config.data.encode_input, encode_output=config.data.encode_output,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8109298a-aca3-45b7-a8de-c5cf4e1c210b",
   "metadata": {},
   "source": [
    "# Creating the model and putting it on the GPU "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "db295d23-ab86-4f37-83cc-7af0a8e485ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Given argument key='dropout' that is not in TFNO2d's signature.\n",
      "Given argument key='tensor_lasso_penalty' that is not in TFNO2d's signature.\n",
      "Keyword argument out_channels not specified for model TFNO2d, using default=1.\n",
      "Keyword argument lifting_channels not specified for model TFNO2d, using default=256.\n",
      "Keyword argument non_linearity not specified for model TFNO2d, using default=<built-in function gelu>.\n",
      "Keyword argument decomposition_kwargs not specified for model TFNO2d, using default={}.\n",
      "\n",
      "n_params: 16844673\n"
     ]
    }
   ],
   "source": [
    "model = get_model(config)\n",
    "model = model.to(device)\n",
    "\n",
    "#Log parameter count\n",
    "if is_logger:\n",
    "    n_params = count_params(model)\n",
    "\n",
    "    if config.verbose:\n",
    "        print(f'\\nn_params: {n_params}')\n",
    "        sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec85d0a-4db4-4b1f-b599-8c2afc98520a",
   "metadata": {},
   "source": [
    "# Create the optimizer and learning rate scheduler\n",
    "\n",
    "Here, we use an Adam optimizer and a learning rate schedule depending on the configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5164537a-267b-4fda-9bcd-257dc3ac4826",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create the optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), \n",
    "                                lr=config.opt.learning_rate, \n",
    "                                weight_decay=config.opt.weight_decay)\n",
    "\n",
    "if config.opt.scheduler == 'ReduceLROnPlateau':\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=config.opt.gamma, patience=config.opt.scheduler_patience, mode='min')\n",
    "elif config.opt.scheduler == 'CosineAnnealingLR':\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.opt.scheduler_T_max)\n",
    "elif config.opt.scheduler == 'StepLR':\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, \n",
    "                                                step_size=config.opt.step_size,\n",
    "                                                gamma=config.opt.gamma)\n",
    "else:\n",
    "    raise ValueError(f'Got {config.opt.scheduler=}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e52a72eb-965a-4997-89a4-0cdfcbcb0a1a",
   "metadata": {},
   "source": [
    "# Creating the loss\n",
    "\n",
    "We will optimize the Sobolev norm but also evaluate our goal: the l2 relative error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "07a53d9d-2d06-4d36-9b46-2c7f15f29c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the losses\n",
    "l2loss = LpLoss(d=2, p=2)\n",
    "h1loss = H1Loss(d=2)\n",
    "if config.opt.training_loss == 'l2':\n",
    "    train_loss = l2loss\n",
    "elif config.opt.training_loss == 'h1':\n",
    "    train_loss = h1loss\n",
    "else:\n",
    "    raise ValueError(f'Got training_loss={config.opt.training_loss} but expected one of [\"l2\", \"h1\"]')\n",
    "eval_losses={'h1': h1loss, 'l2': l2loss}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5dad660e-43e9-4f38-91f6-8427b14b8ae0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "### MODEL ###\n",
      " TFNO2d(\n",
      "  (convs): FactorizedSpectralConv2d(\n",
      "    (weight): ModuleList(\n",
      "      (0): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (1): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (2): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (3): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (4): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (5): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (6): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "      (7): ComplexDenseTensor(shape=torch.Size([64, 64, 16, 16]), rank=None)\n",
      "    )\n",
      "  )\n",
      "  (fno_skips): ModuleList(\n",
      "    (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (3): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "  )\n",
      "  (mlp): ModuleList(\n",
      "    (0): MLP(\n",
      "      (fcs): ModuleList(\n",
      "        (0): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "        (1): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      )\n",
      "    )\n",
      "    (1): MLP(\n",
      "      (fcs): ModuleList(\n",
      "        (0): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "        (1): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      )\n",
      "    )\n",
      "    (2): MLP(\n",
      "      (fcs): ModuleList(\n",
      "        (0): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "        (1): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      )\n",
      "    )\n",
      "    (3): MLP(\n",
      "      (fcs): ModuleList(\n",
      "        (0): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
      "        (1): Conv2d(32, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (mlp_skips): ModuleList(\n",
      "    (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "    (3): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "  )\n",
      "  (norm): ModuleList(\n",
      "    (0): GroupNorm(1, 64, eps=1e-05, affine=True)\n",
      "    (1): GroupNorm(1, 64, eps=1e-05, affine=True)\n",
      "    (2): GroupNorm(1, 64, eps=1e-05, affine=True)\n",
      "    (3): GroupNorm(1, 64, eps=1e-05, affine=True)\n",
      "  )\n",
      "  (lifting): Lifting(\n",
      "    (fc): Conv2d(3, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "  )\n",
      "  (projection): Projection(\n",
      "    (fc1): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (fc2): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))\n",
      "  )\n",
      ")\n",
      "\n",
      "### OPTIMIZER ###\n",
      " Adam (\n",
      "Parameter Group 0\n",
      "    amsgrad: False\n",
      "    betas: (0.9, 0.999)\n",
      "    capturable: False\n",
      "    differentiable: False\n",
      "    eps: 1e-08\n",
      "    foreach: None\n",
      "    fused: False\n",
      "    initial_lr: 0.005\n",
      "    lr: 0.005\n",
      "    maximize: False\n",
      "    weight_decay: 0.0001\n",
      ")\n",
      "\n",
      "### SCHEDULER ###\n",
      " <torch.optim.lr_scheduler.CosineAnnealingLR object at 0x7fff2008b7f0>\n",
      "\n",
      "### LOSSES ###\n",
      "\n",
      " * Train: <neuralop.training.losses.H1Loss object at 0x7fff275a5bb0>\n",
      "\n",
      " * Test: {'h1': <neuralop.training.losses.H1Loss object at 0x7fff275a5bb0>, 'l2': <neuralop.training.losses.LpLoss object at 0x7fff275a5f40>}\n",
      "\n",
      "### Beginning Training...\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if config.verbose and is_logger:\n",
    "    print('\\n### MODEL ###\\n', model)\n",
    "    print('\\n### OPTIMIZER ###\\n', optimizer)\n",
    "    print('\\n### SCHEDULER ###\\n', scheduler)\n",
    "    print('\\n### LOSSES ###')\n",
    "    print(f'\\n * Train: {train_loss}')\n",
    "    print(f'\\n * Test: {eval_losses}')\n",
    "    print(f'\\n### Beginning Training...\\n')\n",
    "    sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5967441-b8bc-4ea8-a4d9-7a5bea384cbf",
   "metadata": {},
   "source": [
    "# Creating the trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a19ebfd3-8a2b-42c0-af98-7a1db2dda0f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on regular inputs (no multi-grid patching).\n",
      "MGPatching(self.n_patches=[1, 1], self.padding_fraction=[0, 0], self.levels=0, use_distributed=False, stitching=False)\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(model, n_epochs=config.opt.n_epochs,\n",
    "                  device=device,\n",
    "                  mg_patching_levels=config.patching.levels,\n",
    "                  mg_patching_padding=config.patching.padding,\n",
    "                  mg_patching_stitching=config.patching.stitching,\n",
    "                  wandb_log=config.wandb.log,\n",
    "                  log_test_interval=config.wandb.log_test_interval,\n",
    "                  log_output=False,\n",
    "                  use_distributed=config.distributed.use_distributed,\n",
    "                  verbose=config.verbose and is_logger)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b16a3727-313d-4219-8f8f-0cec58d74b00",
   "metadata": {},
   "source": [
    "# Training the model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0d6e3298-99ee-4371-8bad-60e6aac03d56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on 3000 samples, testing on [32, 64].\n",
      "[0] time=3.03, avg_loss=7.8899, train_err=0.3945, 32_h1=0.2295, 32_l2=0.1710, 64_h1=0.2847, 64_l2=0.1733\n",
      "[1] time=1.38, avg_loss=3.7664, train_err=0.1883, 32_h1=0.1646, 32_l2=0.1177, 64_h1=0.2326, 64_l2=0.1221\n",
      "[2] time=1.37, avg_loss=3.1005, train_err=0.1550, 32_h1=0.1411, 32_l2=0.1027, 64_h1=0.2156, 64_l2=0.1106\n",
      "[3] time=1.36, avg_loss=2.5222, train_err=0.1261, 32_h1=0.1238, 32_l2=0.0800, 64_h1=0.2026, 64_l2=0.0936\n",
      "[4] time=1.36, avg_loss=2.3043, train_err=0.1152, 32_h1=0.1235, 32_l2=0.0808, 64_h1=0.1874, 64_l2=0.0858\n",
      "[5] time=1.36, avg_loss=2.2108, train_err=0.1105, 32_h1=0.1332, 32_l2=0.1041, 64_h1=0.2055, 64_l2=0.1122\n",
      "[6] time=1.37, avg_loss=1.9753, train_err=0.0988, 32_h1=0.1077, 32_l2=0.0720, 64_h1=0.1885, 64_l2=0.0838\n",
      "[7] time=1.37, avg_loss=1.9352, train_err=0.0968, 32_h1=0.1032, 32_l2=0.0642, 64_h1=0.1847, 64_l2=0.0753\n",
      "[8] time=1.36, avg_loss=1.8174, train_err=0.0909, 32_h1=0.1013, 32_l2=0.0632, 64_h1=0.1798, 64_l2=0.0763\n",
      "[9] time=1.37, avg_loss=1.7847, train_err=0.0892, 32_h1=0.1053, 32_l2=0.0672, 64_h1=0.1909, 64_l2=0.0788\n",
      "[10] time=1.37, avg_loss=1.6375, train_err=0.0819, 32_h1=0.0926, 32_l2=0.0513, 64_h1=0.1808, 64_l2=0.0666\n",
      "[11] time=1.37, avg_loss=1.5826, train_err=0.0791, 32_h1=0.0958, 32_l2=0.0574, 64_h1=0.1810, 64_l2=0.0700\n",
      "[12] time=1.36, avg_loss=1.6231, train_err=0.0812, 32_h1=0.0940, 32_l2=0.0534, 64_h1=0.1740, 64_l2=0.0636\n",
      "[13] time=1.42, avg_loss=1.5427, train_err=0.0771, 32_h1=0.0937, 32_l2=0.0532, 64_h1=0.1834, 64_l2=0.0692\n",
      "[14] time=1.37, avg_loss=1.4741, train_err=0.0737, 32_h1=0.0989, 32_l2=0.0623, 64_h1=0.1844, 64_l2=0.0798\n",
      "[15] time=1.36, avg_loss=1.5156, train_err=0.0758, 32_h1=0.1020, 32_l2=0.0649, 64_h1=0.1844, 64_l2=0.0730\n",
      "[16] time=1.37, avg_loss=1.5620, train_err=0.0781, 32_h1=0.0940, 32_l2=0.0608, 64_h1=0.1803, 64_l2=0.0747\n",
      "[17] time=1.36, avg_loss=1.3939, train_err=0.0697, 32_h1=0.1018, 32_l2=0.0620, 64_h1=0.1842, 64_l2=0.0772\n",
      "[18] time=1.89, avg_loss=1.4904, train_err=0.0745, 32_h1=0.1010, 32_l2=0.0704, 64_h1=0.1868, 64_l2=0.0794\n",
      "[19] time=1.83, avg_loss=1.4300, train_err=0.0715, 32_h1=0.0929, 32_l2=0.0525, 64_h1=0.1784, 64_l2=0.0679\n",
      "[20] time=1.84, avg_loss=1.3752, train_err=0.0688, 32_h1=0.0964, 32_l2=0.0635, 64_h1=0.1825, 64_l2=0.0694\n",
      "[21] time=1.84, avg_loss=1.4671, train_err=0.0734, 32_h1=0.0911, 32_l2=0.0513, 64_h1=0.1832, 64_l2=0.0696\n",
      "[22] time=1.88, avg_loss=1.3043, train_err=0.0652, 32_h1=0.0938, 32_l2=0.0538, 64_h1=0.1804, 64_l2=0.0687\n",
      "[23] time=1.37, avg_loss=1.2880, train_err=0.0644, 32_h1=0.0897, 32_l2=0.0492, 64_h1=0.1824, 64_l2=0.0629\n",
      "[24] time=1.37, avg_loss=1.3901, train_err=0.0695, 32_h1=0.1080, 32_l2=0.0701, 64_h1=0.1828, 64_l2=0.0785\n",
      "[25] time=1.37, avg_loss=1.3788, train_err=0.0689, 32_h1=0.0878, 32_l2=0.0514, 64_h1=0.1744, 64_l2=0.0613\n",
      "[26] time=1.37, avg_loss=1.3071, train_err=0.0654, 32_h1=0.0880, 32_l2=0.0489, 64_h1=0.1847, 64_l2=0.0698\n",
      "[27] time=1.36, avg_loss=1.3056, train_err=0.0653, 32_h1=0.0980, 32_l2=0.0679, 64_h1=0.1828, 64_l2=0.0830\n",
      "[28] time=1.37, avg_loss=1.2677, train_err=0.0634, 32_h1=0.0956, 32_l2=0.0621, 64_h1=0.1827, 64_l2=0.0692\n",
      "[29] time=1.37, avg_loss=1.2611, train_err=0.0631, 32_h1=0.0913, 32_l2=0.0500, 64_h1=0.1855, 64_l2=0.0652\n",
      "[30] time=1.37, avg_loss=1.1833, train_err=0.0592, 32_h1=0.0888, 32_l2=0.0512, 64_h1=0.1818, 64_l2=0.0655\n",
      "[31] time=1.36, avg_loss=1.2170, train_err=0.0608, 32_h1=0.0879, 32_l2=0.0481, 64_h1=0.1758, 64_l2=0.0625\n",
      "[32] time=1.36, avg_loss=1.1431, train_err=0.0572, 32_h1=0.0886, 32_l2=0.0479, 64_h1=0.1756, 64_l2=0.0594\n",
      "[33] time=1.37, avg_loss=1.2162, train_err=0.0608, 32_h1=0.0923, 32_l2=0.0522, 64_h1=0.1749, 64_l2=0.0629\n",
      "[34] time=1.37, avg_loss=1.1588, train_err=0.0579, 32_h1=0.0892, 32_l2=0.0526, 64_h1=0.1797, 64_l2=0.0656\n",
      "[35] time=1.37, avg_loss=1.1747, train_err=0.0587, 32_h1=0.0884, 32_l2=0.0481, 64_h1=0.1829, 64_l2=0.0650\n",
      "[36] time=1.36, avg_loss=1.1491, train_err=0.0575, 32_h1=0.0936, 32_l2=0.0542, 64_h1=0.1787, 64_l2=0.0672\n",
      "[37] time=1.37, avg_loss=1.1532, train_err=0.0577, 32_h1=0.0950, 32_l2=0.0569, 64_h1=0.1737, 64_l2=0.0679\n",
      "[38] time=1.37, avg_loss=1.2426, train_err=0.0621, 32_h1=0.0875, 32_l2=0.0488, 64_h1=0.1750, 64_l2=0.0638\n",
      "[39] time=1.37, avg_loss=1.1345, train_err=0.0567, 32_h1=0.0874, 32_l2=0.0493, 64_h1=0.1780, 64_l2=0.0658\n",
      "[40] time=1.36, avg_loss=1.1238, train_err=0.0562, 32_h1=0.0914, 32_l2=0.0516, 64_h1=0.1796, 64_l2=0.0662\n",
      "[41] time=1.36, avg_loss=1.1093, train_err=0.0555, 32_h1=0.0855, 32_l2=0.0457, 64_h1=0.1741, 64_l2=0.0621\n",
      "[42] time=1.36, avg_loss=1.0772, train_err=0.0539, 32_h1=0.0899, 32_l2=0.0523, 64_h1=0.1807, 64_l2=0.0688\n",
      "[43] time=1.36, avg_loss=1.0772, train_err=0.0539, 32_h1=0.0894, 32_l2=0.0556, 64_h1=0.1769, 64_l2=0.0705\n",
      "[44] time=1.36, avg_loss=1.0901, train_err=0.0545, 32_h1=0.0843, 32_l2=0.0443, 64_h1=0.1750, 64_l2=0.0589\n",
      "[45] time=1.36, avg_loss=1.0783, train_err=0.0539, 32_h1=0.0874, 32_l2=0.0486, 64_h1=0.1778, 64_l2=0.0593\n",
      "[46] time=1.46, avg_loss=1.0837, train_err=0.0542, 32_h1=0.0874, 32_l2=0.0482, 64_h1=0.1722, 64_l2=0.0575\n",
      "[47] time=1.37, avg_loss=1.1760, train_err=0.0588, 32_h1=0.0873, 32_l2=0.0507, 64_h1=0.1706, 64_l2=0.0639\n",
      "[48] time=1.37, avg_loss=1.0357, train_err=0.0518, 32_h1=0.0889, 32_l2=0.0503, 64_h1=0.1799, 64_l2=0.0663\n",
      "[49] time=1.36, avg_loss=1.0873, train_err=0.0544, 32_h1=0.0846, 32_l2=0.0464, 64_h1=0.1725, 64_l2=0.0592\n",
      "[50] time=1.36, avg_loss=1.0996, train_err=0.0550, 32_h1=0.0861, 32_l2=0.0461, 64_h1=0.1696, 64_l2=0.0598\n",
      "[51] time=1.37, avg_loss=1.0487, train_err=0.0524, 32_h1=0.0839, 32_l2=0.0433, 64_h1=0.1752, 64_l2=0.0602\n",
      "[52] time=1.37, avg_loss=1.0527, train_err=0.0526, 32_h1=0.0858, 32_l2=0.0469, 64_h1=0.1736, 64_l2=0.0588\n",
      "[53] time=1.36, avg_loss=1.0138, train_err=0.0507, 32_h1=0.0854, 32_l2=0.0475, 64_h1=0.1777, 64_l2=0.0619\n",
      "[54] time=1.36, avg_loss=1.0210, train_err=0.0511, 32_h1=0.0832, 32_l2=0.0431, 64_h1=0.1728, 64_l2=0.0580\n",
      "[55] time=1.36, avg_loss=0.9939, train_err=0.0497, 32_h1=0.0870, 32_l2=0.0474, 64_h1=0.1755, 64_l2=0.0609\n",
      "[56] time=1.37, avg_loss=1.0085, train_err=0.0504, 32_h1=0.0833, 32_l2=0.0438, 64_h1=0.1731, 64_l2=0.0603\n",
      "[57] time=1.37, avg_loss=1.0132, train_err=0.0507, 32_h1=0.0842, 32_l2=0.0462, 64_h1=0.1757, 64_l2=0.0613\n",
      "[58] time=1.36, avg_loss=0.9938, train_err=0.0497, 32_h1=0.0839, 32_l2=0.0439, 64_h1=0.1811, 64_l2=0.0651\n",
      "[59] time=1.36, avg_loss=0.9814, train_err=0.0491, 32_h1=0.0820, 32_l2=0.0425, 64_h1=0.1728, 64_l2=0.0565\n",
      "[60] time=1.36, avg_loss=0.9849, train_err=0.0492, 32_h1=0.0861, 32_l2=0.0477, 64_h1=0.1715, 64_l2=0.0616\n",
      "[61] time=1.37, avg_loss=0.9787, train_err=0.0489, 32_h1=0.0844, 32_l2=0.0450, 64_h1=0.1742, 64_l2=0.0623\n",
      "[62] time=1.36, avg_loss=1.0104, train_err=0.0505, 32_h1=0.0830, 32_l2=0.0437, 64_h1=0.1769, 64_l2=0.0605\n",
      "[63] time=1.36, avg_loss=0.9910, train_err=0.0495, 32_h1=0.0821, 32_l2=0.0415, 64_h1=0.1742, 64_l2=0.0579\n",
      "[64] time=1.36, avg_loss=0.9622, train_err=0.0481, 32_h1=0.0849, 32_l2=0.0462, 64_h1=0.1763, 64_l2=0.0608\n",
      "[65] time=1.36, avg_loss=1.0191, train_err=0.0510, 32_h1=0.0823, 32_l2=0.0419, 64_h1=0.1736, 64_l2=0.0570\n",
      "[66] time=1.37, avg_loss=0.9814, train_err=0.0491, 32_h1=0.0873, 32_l2=0.0492, 64_h1=0.1752, 64_l2=0.0643\n",
      "[67] time=1.36, avg_loss=0.9867, train_err=0.0493, 32_h1=0.0833, 32_l2=0.0446, 64_h1=0.1698, 64_l2=0.0588\n",
      "[68] time=1.36, avg_loss=0.9983, train_err=0.0499, 32_h1=0.0815, 32_l2=0.0417, 64_h1=0.1712, 64_l2=0.0590\n",
      "[69] time=1.37, avg_loss=0.9956, train_err=0.0498, 32_h1=0.0836, 32_l2=0.0453, 64_h1=0.1756, 64_l2=0.0604\n",
      "[70] time=1.37, avg_loss=0.9433, train_err=0.0472, 32_h1=0.0830, 32_l2=0.0432, 64_h1=0.1739, 64_l2=0.0583\n",
      "[71] time=1.36, avg_loss=0.9813, train_err=0.0491, 32_h1=0.0830, 32_l2=0.0433, 64_h1=0.1691, 64_l2=0.0588\n",
      "[72] time=1.36, avg_loss=0.9456, train_err=0.0473, 32_h1=0.0828, 32_l2=0.0429, 64_h1=0.1695, 64_l2=0.0599\n",
      "[73] time=1.37, avg_loss=0.9099, train_err=0.0455, 32_h1=0.0835, 32_l2=0.0438, 64_h1=0.1716, 64_l2=0.0599\n",
      "[74] time=1.37, avg_loss=0.9241, train_err=0.0462, 32_h1=0.0816, 32_l2=0.0419, 64_h1=0.1699, 64_l2=0.0572\n",
      "[75] time=1.37, avg_loss=0.8907, train_err=0.0445, 32_h1=0.0825, 32_l2=0.0410, 64_h1=0.1772, 64_l2=0.0604\n",
      "[76] time=1.36, avg_loss=0.8940, train_err=0.0447, 32_h1=0.0821, 32_l2=0.0428, 64_h1=0.1733, 64_l2=0.0588\n",
      "[77] time=1.37, avg_loss=0.8958, train_err=0.0448, 32_h1=0.0828, 32_l2=0.0447, 64_h1=0.1756, 64_l2=0.0593\n",
      "[78] time=1.37, avg_loss=0.9276, train_err=0.0464, 32_h1=0.0816, 32_l2=0.0424, 64_h1=0.1740, 64_l2=0.0599\n",
      "[79] time=1.37, avg_loss=0.8763, train_err=0.0438, 32_h1=0.0818, 32_l2=0.0414, 64_h1=0.1715, 64_l2=0.0570\n",
      "[80] time=1.36, avg_loss=0.8634, train_err=0.0432, 32_h1=0.0812, 32_l2=0.0416, 64_h1=0.1753, 64_l2=0.0614\n",
      "[81] time=1.36, avg_loss=0.8450, train_err=0.0423, 32_h1=0.0832, 32_l2=0.0448, 64_h1=0.1701, 64_l2=0.0626\n",
      "[82] time=1.37, avg_loss=0.8997, train_err=0.0450, 32_h1=0.0818, 32_l2=0.0419, 64_h1=0.1718, 64_l2=0.0590\n",
      "[83] time=1.37, avg_loss=0.8658, train_err=0.0433, 32_h1=0.0816, 32_l2=0.0415, 64_h1=0.1703, 64_l2=0.0552\n",
      "[84] time=1.37, avg_loss=0.9292, train_err=0.0465, 32_h1=0.0815, 32_l2=0.0424, 64_h1=0.1674, 64_l2=0.0580\n",
      "[85] time=1.36, avg_loss=0.9417, train_err=0.0471, 32_h1=0.0825, 32_l2=0.0439, 64_h1=0.1755, 64_l2=0.0608\n",
      "[86] time=1.37, avg_loss=0.8608, train_err=0.0430, 32_h1=0.0792, 32_l2=0.0392, 64_h1=0.1720, 64_l2=0.0573\n",
      "[87] time=1.38, avg_loss=0.9083, train_err=0.0454, 32_h1=0.0822, 32_l2=0.0440, 64_h1=0.1693, 64_l2=0.0602\n",
      "[88] time=1.57, avg_loss=0.8522, train_err=0.0426, 32_h1=0.0823, 32_l2=0.0427, 64_h1=0.1695, 64_l2=0.0571\n",
      "[89] time=1.36, avg_loss=0.8273, train_err=0.0414, 32_h1=0.0813, 32_l2=0.0414, 64_h1=0.1702, 64_l2=0.0568\n",
      "[90] time=1.36, avg_loss=0.8612, train_err=0.0431, 32_h1=0.0834, 32_l2=0.0468, 64_h1=0.1718, 64_l2=0.0641\n",
      "[91] time=1.36, avg_loss=0.8358, train_err=0.0418, 32_h1=0.0811, 32_l2=0.0410, 64_h1=0.1678, 64_l2=0.0558\n",
      "[92] time=1.37, avg_loss=0.8725, train_err=0.0436, 32_h1=0.0807, 32_l2=0.0408, 64_h1=0.1688, 64_l2=0.0557\n",
      "[93] time=1.36, avg_loss=0.8163, train_err=0.0408, 32_h1=0.0804, 32_l2=0.0417, 64_h1=0.1714, 64_l2=0.0593\n",
      "[94] time=1.36, avg_loss=0.8119, train_err=0.0406, 32_h1=0.0791, 32_l2=0.0393, 64_h1=0.1706, 64_l2=0.0581\n",
      "[95] time=1.36, avg_loss=0.8022, train_err=0.0401, 32_h1=0.0819, 32_l2=0.0416, 64_h1=0.1697, 64_l2=0.0555\n",
      "[96] time=1.37, avg_loss=0.8371, train_err=0.0419, 32_h1=0.0793, 32_l2=0.0393, 64_h1=0.1684, 64_l2=0.0570\n",
      "[97] time=1.37, avg_loss=0.8227, train_err=0.0411, 32_h1=0.0800, 32_l2=0.0407, 64_h1=0.1685, 64_l2=0.0583\n",
      "[98] time=1.43, avg_loss=0.8176, train_err=0.0409, 32_h1=0.0841, 32_l2=0.0471, 64_h1=0.1681, 64_l2=0.0578\n",
      "[99] time=1.85, avg_loss=0.8517, train_err=0.0426, 32_h1=0.0809, 32_l2=0.0401, 64_h1=0.1726, 64_l2=0.0607\n",
      "[100] time=1.85, avg_loss=0.8445, train_err=0.0422, 32_h1=0.0810, 32_l2=0.0408, 64_h1=0.1688, 64_l2=0.0558\n",
      "[101] time=1.42, avg_loss=0.7962, train_err=0.0398, 32_h1=0.0796, 32_l2=0.0393, 64_h1=0.1680, 64_l2=0.0577\n",
      "[102] time=1.84, avg_loss=0.7758, train_err=0.0388, 32_h1=0.0799, 32_l2=0.0398, 64_h1=0.1664, 64_l2=0.0556\n",
      "[103] time=1.87, avg_loss=0.8005, train_err=0.0400, 32_h1=0.0792, 32_l2=0.0395, 64_h1=0.1688, 64_l2=0.0552\n",
      "[104] time=1.43, avg_loss=0.8099, train_err=0.0405, 32_h1=0.0791, 32_l2=0.0394, 64_h1=0.1664, 64_l2=0.0535\n",
      "[105] time=1.37, avg_loss=0.7828, train_err=0.0391, 32_h1=0.0815, 32_l2=0.0430, 64_h1=0.1691, 64_l2=0.0574\n",
      "[106] time=1.37, avg_loss=0.7799, train_err=0.0390, 32_h1=0.0795, 32_l2=0.0393, 64_h1=0.1679, 64_l2=0.0556\n",
      "[107] time=1.36, avg_loss=0.7685, train_err=0.0384, 32_h1=0.0810, 32_l2=0.0434, 64_h1=0.1725, 64_l2=0.0633\n",
      "[108] time=1.36, avg_loss=0.7581, train_err=0.0379, 32_h1=0.0801, 32_l2=0.0407, 64_h1=0.1744, 64_l2=0.0574\n",
      "[109] time=1.37, avg_loss=0.7415, train_err=0.0371, 32_h1=0.0782, 32_l2=0.0383, 64_h1=0.1670, 64_l2=0.0540\n",
      "[110] time=1.37, avg_loss=0.7387, train_err=0.0369, 32_h1=0.0790, 32_l2=0.0392, 64_h1=0.1664, 64_l2=0.0539\n",
      "[111] time=1.37, avg_loss=0.7338, train_err=0.0367, 32_h1=0.0788, 32_l2=0.0385, 64_h1=0.1694, 64_l2=0.0574\n",
      "[112] time=1.36, avg_loss=0.7426, train_err=0.0371, 32_h1=0.0811, 32_l2=0.0434, 64_h1=0.1745, 64_l2=0.0593\n",
      "[113] time=1.36, avg_loss=0.7849, train_err=0.0392, 32_h1=0.0817, 32_l2=0.0452, 64_h1=0.1653, 64_l2=0.0627\n",
      "[114] time=1.37, avg_loss=0.7933, train_err=0.0397, 32_h1=0.0803, 32_l2=0.0409, 64_h1=0.1715, 64_l2=0.0568\n",
      "[115] time=1.37, avg_loss=0.7377, train_err=0.0369, 32_h1=0.0789, 32_l2=0.0389, 64_h1=0.1688, 64_l2=0.0556\n",
      "[116] time=1.37, avg_loss=0.7639, train_err=0.0382, 32_h1=0.0794, 32_l2=0.0394, 64_h1=0.1683, 64_l2=0.0574\n",
      "[117] time=1.36, avg_loss=0.7515, train_err=0.0376, 32_h1=0.0785, 32_l2=0.0382, 64_h1=0.1665, 64_l2=0.0549\n",
      "[118] time=1.37, avg_loss=0.7180, train_err=0.0359, 32_h1=0.0792, 32_l2=0.0394, 64_h1=0.1671, 64_l2=0.0576\n",
      "[119] time=1.37, avg_loss=0.7191, train_err=0.0360, 32_h1=0.0795, 32_l2=0.0396, 64_h1=0.1672, 64_l2=0.0541\n",
      "[120] time=1.37, avg_loss=0.7148, train_err=0.0357, 32_h1=0.0792, 32_l2=0.0389, 64_h1=0.1671, 64_l2=0.0575\n",
      "[121] time=1.36, avg_loss=0.7012, train_err=0.0351, 32_h1=0.0795, 32_l2=0.0399, 64_h1=0.1639, 64_l2=0.0555\n",
      "[122] time=1.37, avg_loss=0.6962, train_err=0.0348, 32_h1=0.0787, 32_l2=0.0388, 64_h1=0.1697, 64_l2=0.0570\n",
      "[123] time=1.37, avg_loss=0.6970, train_err=0.0349, 32_h1=0.0793, 32_l2=0.0388, 64_h1=0.1693, 64_l2=0.0567\n",
      "[124] time=1.37, avg_loss=0.6888, train_err=0.0344, 32_h1=0.0788, 32_l2=0.0382, 64_h1=0.1687, 64_l2=0.0570\n",
      "[125] time=1.37, avg_loss=0.7060, train_err=0.0353, 32_h1=0.0799, 32_l2=0.0412, 64_h1=0.1649, 64_l2=0.0576\n",
      "[126] time=1.36, avg_loss=0.6991, train_err=0.0350, 32_h1=0.0792, 32_l2=0.0393, 64_h1=0.1681, 64_l2=0.0583\n",
      "[127] time=1.37, avg_loss=0.7098, train_err=0.0355, 32_h1=0.0796, 32_l2=0.0406, 64_h1=0.1641, 64_l2=0.0574\n",
      "[128] time=1.37, avg_loss=0.6971, train_err=0.0349, 32_h1=0.0792, 32_l2=0.0399, 64_h1=0.1690, 64_l2=0.0588\n",
      "[129] time=1.37, avg_loss=0.6810, train_err=0.0340, 32_h1=0.0793, 32_l2=0.0393, 64_h1=0.1648, 64_l2=0.0559\n",
      "[130] time=1.36, avg_loss=0.6848, train_err=0.0342, 32_h1=0.0780, 32_l2=0.0378, 64_h1=0.1670, 64_l2=0.0536\n",
      "[131] time=1.94, avg_loss=0.6600, train_err=0.0330, 32_h1=0.0779, 32_l2=0.0379, 64_h1=0.1661, 64_l2=0.0545\n",
      "[132] time=1.87, avg_loss=0.6428, train_err=0.0321, 32_h1=0.0794, 32_l2=0.0394, 64_h1=0.1695, 64_l2=0.0588\n",
      "[133] time=1.85, avg_loss=0.6532, train_err=0.0327, 32_h1=0.0789, 32_l2=0.0392, 64_h1=0.1690, 64_l2=0.0568\n",
      "[134] time=1.88, avg_loss=0.6573, train_err=0.0329, 32_h1=0.0780, 32_l2=0.0376, 64_h1=0.1664, 64_l2=0.0559\n",
      "[135] time=1.78, avg_loss=0.6445, train_err=0.0322, 32_h1=0.0784, 32_l2=0.0386, 64_h1=0.1644, 64_l2=0.0560\n",
      "[136] time=1.82, avg_loss=0.6378, train_err=0.0319, 32_h1=0.0780, 32_l2=0.0383, 64_h1=0.1651, 64_l2=0.0527\n",
      "[137] time=1.85, avg_loss=0.6550, train_err=0.0327, 32_h1=0.0792, 32_l2=0.0400, 64_h1=0.1652, 64_l2=0.0551\n",
      "[138] time=1.75, avg_loss=0.6341, train_err=0.0317, 32_h1=0.0775, 32_l2=0.0376, 64_h1=0.1661, 64_l2=0.0556\n",
      "[139] time=2.00, avg_loss=0.8234, train_err=0.0412, 32_h1=0.0783, 32_l2=0.0386, 64_h1=0.1654, 64_l2=0.0566\n",
      "[140] time=1.97, avg_loss=0.6822, train_err=0.0341, 32_h1=0.0784, 32_l2=0.0380, 64_h1=0.1675, 64_l2=0.0558\n",
      "[141] time=1.98, avg_loss=0.6332, train_err=0.0317, 32_h1=0.0778, 32_l2=0.0379, 64_h1=0.1670, 64_l2=0.0546\n",
      "[142] time=1.99, avg_loss=0.6205, train_err=0.0310, 32_h1=0.0786, 32_l2=0.0394, 64_h1=0.1678, 64_l2=0.0592\n",
      "[143] time=1.82, avg_loss=0.6098, train_err=0.0305, 32_h1=0.0785, 32_l2=0.0386, 64_h1=0.1676, 64_l2=0.0584\n",
      "[144] time=1.38, avg_loss=0.6116, train_err=0.0306, 32_h1=0.0794, 32_l2=0.0422, 64_h1=0.1702, 64_l2=0.0591\n",
      "[145] time=1.37, avg_loss=0.6018, train_err=0.0301, 32_h1=0.0776, 32_l2=0.0373, 64_h1=0.1674, 64_l2=0.0564\n",
      "[146] time=1.38, avg_loss=0.6001, train_err=0.0300, 32_h1=0.0781, 32_l2=0.0386, 64_h1=0.1662, 64_l2=0.0583\n",
      "[147] time=1.38, avg_loss=0.5990, train_err=0.0300, 32_h1=0.0796, 32_l2=0.0416, 64_h1=0.1679, 64_l2=0.0572\n",
      "[148] time=1.38, avg_loss=0.6462, train_err=0.0323, 32_h1=0.0802, 32_l2=0.0411, 64_h1=0.1721, 64_l2=0.0580\n",
      "[149] time=1.37, avg_loss=0.6152, train_err=0.0308, 32_h1=0.0777, 32_l2=0.0373, 64_h1=0.1688, 64_l2=0.0562\n"
     ]
    }
   ],
   "source": [
    "trainer.train(train_loader, test_loaders,\n",
    "              output_encoder,\n",
    "              model, \n",
    "              optimizer,\n",
    "              scheduler, \n",
    "              regularizer=False, \n",
    "              training_loss=train_loss,\n",
    "              eval_losses=eval_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b20be56-d200-44dc-b97b-fca021e353c8",
   "metadata": {},
   "source": [
    "# Follow-up questions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a67e1d5-4b9a-4be3-bff4-fb2a6b152f9c",
   "metadata": {},
   "source": [
    "You can now play with the configuration and see how the performance is impacted.\n",
    "\n",
    "Which parameters do you think will most influence performance? \n",
    "Learning rate? Learning schedule? hidden_channels? Number of training samples? \n",
    "\n",
    "Does your intuition match the results you are getting?"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}