fno_task.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fourier Neural Operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "torch.manual_seed(2002)\n",
    "np.random.seed(2002)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Постановка задачи\n",
    "\n",
    "В качестве входных данных сеть получает начальное условие для уравнения Бюргерса $$ {\\frac {\\partial u}{\\partial t}}+u{\\frac {\\partial u}{\\partial x}}=\\nu {\\frac {\\partial ^{2}u}{\\partial x^{2}}}$$ на интервале $[0, 0.5]$ и экстраполирует его на интервал $[0.5, 1]$. Размер равномерной сетки - 256 элементов."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Этап 1: обучение FNO из пакета neuralop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "На этом этапе задача будет решаться с официальной реализацией FNO из пакета neuralop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Загрузка и обработка данных"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Путь к данным. Лучше поменять, так как я прописывал его относительно базовой директории окружения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets_path = \"./Task2_with_datasets/datasets\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataloaders.burgers import load_burgers_1dtime\n",
    "from dataloaders.tensor_dataset import TensorDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "При загрузке данным надо вытащить сами массивы данных, чтобы разбить их на x и y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, _, train, test = load_burgers_1dtime(\n",
    "    data_path=os.path.join(datasets_path, \"burgers\"),\n",
    "    n_train=400, n_test=100, batch_size=16, batch_size_test=16,\n",
    "    temporal_length=1, spatial_length=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([400, 3, 1, 256]) torch.Size([100, 3, 1, 256])\n"
     ]
    }
   ],
   "source": [
    "print(train.shape, test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([400, 3, 256]) torch.Size([100, 3, 256])\n"
     ]
    }
   ],
   "source": [
    "train = torch.squeeze(train, dim=2)\n",
    "test = torch.squeeze(test, dim=2)\n",
    "print(train.shape, test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = train[:, :, :128]\n",
    "y_train =  torch.unsqueeze(train[:, 0, 128:], dim=1)\n",
    "\n",
    "x_test = test[:, :, :128]\n",
    "y_test = torch.unsqueeze(test[:, 0, 128:], dim=1)\n",
    "train_db = TensorDataset(x_train, y_train)\n",
    "train_loader = torch.utils.data.DataLoader(train_db, batch_size=32, shuffle=False)\n",
    "\n",
    "test_db = TensorDataset(x_test, y_test)\n",
    "test_loader = torch.utils.data.DataLoader(test_db, batch_size=16, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Создание модели"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuralop.models import FNO\n",
    "from neuralop.utils import count_params\n",
    "from neuralop import LpLoss, H1Loss\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Создадим 1-d FNO:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Our model has 205249 parameters.\n",
      "\n",
      "### MODEL ###\n",
      " FNO(\n",
      "  (fno_blocks): FNOBlocks(\n",
      "    (convs): FactorizedSpectralConv(\n",
      "      (weight): ModuleList(\n",
      "        (0-5): 6 x ComplexDenseTensor(shape=torch.Size([32, 32, 16]), rank=None)\n",
      "      )\n",
      "    )\n",
      "    (fno_skips): ModuleList(\n",
      "      (0-5): 6 x Conv1d(32, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
      "    )\n",
      "  )\n",
      "  (lifting): Lifting(\n",
      "    (fc): Conv1d(3, 32, kernel_size=(1,), stride=(1,))\n",
      "  )\n",
      "  (projection): Projection(\n",
      "    (fc1): Conv1d(32, 64, kernel_size=(1,), stride=(1,))\n",
      "    (fc2): Conv1d(64, 1, kernel_size=(1,), stride=(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: None\n",
      "    initial_lr: 0.001\n",
      "    lr: 0.001\n",
      "    maximize: False\n",
      "    weight_decay: 0\n",
      ")\n",
      "\n",
      "### SCHEDULER ###\n",
      " <torch.optim.lr_scheduler.StepLR object at 0x7f53c2307f10>\n",
      "\n",
      "### LOSSES ###\n",
      "\n",
      " * Train: <neuralop.training.losses.H1Loss object at 0x7f54a0086bb0>\n",
      "\n",
      " * Test: {'h1': <neuralop.training.losses.H1Loss object at 0x7f54a0086bb0>, 'l2': <neuralop.training.losses.LpLoss object at 0x7f533c08f1f0>}\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\"\n",
    "\n",
    "model = FNO(n_modes=(32,), in_channels=3, out_channels=1, n_layers=6, hidden_channels=32, projection_channels=64)\n",
    "model = model.to(device)\n",
    "\n",
    "n_params = count_params(model)\n",
    "print(f'\\nOur model has {n_params} parameters.')\n",
    "sys.stdout.flush()\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(),\n",
    "                             lr=1e-3)\n",
    "\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10000, gamma=1)\n",
    "# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=5e-4)\n",
    "\n",
    "h1loss = H1Loss(d=1, reductions='mean')\n",
    "l2loss =LpLoss(d=1, p=2, reductions='mean')\n",
    "\n",
    "train_loss = h1loss\n",
    "eval_losses={'h1': h1loss, 'l2': l2loss}\n",
    "\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",
    "sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Создаем Trainer с коллбэком для Tensorboard и обучаем модель:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-12-02 14:22:07.327745: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-02 14:22:07.759744: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "from callbacks.tensorboard_callback import TensorBoardCallback\n",
    "from callbacks.trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_num = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using standard method to load data to device.\n",
      "using standard method to compute loss.\n",
      "Training on regular inputs (no multi-grid patching).\n",
      "Training on 400 samples\n",
      "Testing on [100] samples         on resolutions ['test'].\n",
      "Training on raw inputs of size x.shape=torch.Size([32, 3, 128]), y.shape=torch.Size([32, 1, 128])\n",
      ".. patched inputs of size x.shape=torch.Size([32, 3, 128]), y.shape=torch.Size([32, 1, 128])\n",
      "Raw outputs of size out.shape=torch.Size([32, 1, 128])\n",
      ".. Processed (unpatched) outputs of size out.shape=torch.Size([32, 1, 128])\n",
      "[0] time=0.16, avg_loss=2.0446, train_err=1.0325, test_h1=1.0377, test_l2=1.0858\n",
      "[3] time=0.12, avg_loss=1.9240, train_err=0.9716, test_h1=0.9322, test_l2=0.9257\n",
      "[6] time=0.12, avg_loss=1.6995, train_err=0.8582, test_h1=0.7847, test_l2=0.8255\n",
      "[9] time=0.12, avg_loss=1.5815, train_err=0.7986, test_h1=0.7334, test_l2=0.8382\n",
      "[12] time=0.12, avg_loss=1.5290, train_err=0.7722, test_h1=0.6954, test_l2=0.7797\n",
      "[15] time=0.12, avg_loss=1.3320, train_err=0.6726, test_h1=0.6219, test_l2=0.6436\n",
      "[18] time=0.12, avg_loss=1.2363, train_err=0.6243, test_h1=0.5585, test_l2=0.5684\n",
      "[21] time=0.12, avg_loss=1.1955, train_err=0.6037, test_h1=0.5439, test_l2=0.5862\n",
      "[24] time=0.12, avg_loss=1.1889, train_err=0.6004, test_h1=0.5482, test_l2=0.6017\n",
      "[27] time=0.12, avg_loss=1.1792, train_err=0.5955, test_h1=0.5487, test_l2=0.6083\n",
      "[30] time=0.12, avg_loss=1.1756, train_err=0.5937, test_h1=0.5495, test_l2=0.6116\n",
      "[33] time=0.12, avg_loss=1.1728, train_err=0.5923, test_h1=0.5492, test_l2=0.6101\n",
      "[36] time=0.12, avg_loss=1.1700, train_err=0.5909, test_h1=0.5448, test_l2=0.5964\n",
      "[39] time=0.12, avg_loss=1.1674, train_err=0.5896, test_h1=0.5387, test_l2=0.5745\n",
      "[42] time=0.12, avg_loss=1.1700, train_err=0.5908, test_h1=0.5376, test_l2=0.5521\n",
      "[45] time=0.12, avg_loss=1.1742, train_err=0.5930, test_h1=0.5529, test_l2=0.5504\n",
      "[48] time=0.12, avg_loss=1.1656, train_err=0.5886, test_h1=0.5499, test_l2=0.5478\n",
      "[51] time=0.12, avg_loss=1.1590, train_err=0.5853, test_h1=0.5475, test_l2=0.5469\n",
      "[54] time=0.12, avg_loss=1.1551, train_err=0.5833, test_h1=0.5486, test_l2=0.5487\n",
      "[57] time=0.12, avg_loss=1.1502, train_err=0.5809, test_h1=0.5492, test_l2=0.5512\n",
      "[60] time=0.12, avg_loss=1.1451, train_err=0.5783, test_h1=0.5501, test_l2=0.5551\n",
      "[63] time=0.12, avg_loss=1.1402, train_err=0.5758, test_h1=0.5515, test_l2=0.5605\n",
      "[66] time=0.12, avg_loss=1.1352, train_err=0.5733, test_h1=0.5530, test_l2=0.5659\n",
      "[69] time=0.12, avg_loss=1.1298, train_err=0.5705, test_h1=0.5535, test_l2=0.5703\n",
      "[72] time=0.12, avg_loss=1.1243, train_err=0.5678, test_h1=0.5535, test_l2=0.5736\n",
      "[75] time=0.12, avg_loss=1.1189, train_err=0.5650, test_h1=0.5531, test_l2=0.5755\n",
      "[78] time=0.12, avg_loss=1.1135, train_err=0.5623, test_h1=0.5526, test_l2=0.5766\n",
      "[81] time=0.12, avg_loss=1.1080, train_err=0.5596, test_h1=0.5521, test_l2=0.5774\n",
      "[84] time=0.12, avg_loss=1.1027, train_err=0.5569, test_h1=0.5526, test_l2=0.5775\n",
      "[87] time=0.12, avg_loss=1.0969, train_err=0.5539, test_h1=0.5540, test_l2=0.5789\n",
      "[90] time=0.12, avg_loss=1.0897, train_err=0.5503, test_h1=0.5534, test_l2=0.5773\n",
      "[93] time=0.12, avg_loss=1.0833, train_err=0.5471, test_h1=0.5529, test_l2=0.5739\n",
      "[96] time=0.12, avg_loss=1.0777, train_err=0.5442, test_h1=0.5541, test_l2=0.5740\n",
      "[99] time=0.12, avg_loss=1.0728, train_err=0.5417, test_h1=0.5558, test_l2=0.5741\n",
      "[102] time=0.12, avg_loss=1.0695, train_err=0.5401, test_h1=0.5574, test_l2=0.5659\n",
      "[105] time=0.12, avg_loss=1.0714, train_err=0.5411, test_h1=0.5606, test_l2=0.5543\n",
      "[108] time=0.12, avg_loss=1.0671, train_err=0.5389, test_h1=0.5573, test_l2=0.5602\n",
      "[111] time=0.12, avg_loss=1.0600, train_err=0.5353, test_h1=0.5530, test_l2=0.5604\n",
      "[114] time=0.12, avg_loss=1.0526, train_err=0.5316, test_h1=0.5530, test_l2=0.5610\n",
      "[117] time=0.12, avg_loss=1.0486, train_err=0.5296, test_h1=0.5540, test_l2=0.5597\n",
      "[120] time=0.12, avg_loss=1.0458, train_err=0.5281, test_h1=0.5536, test_l2=0.5580\n",
      "[123] time=0.12, avg_loss=1.0410, train_err=0.5257, test_h1=0.5531, test_l2=0.5587\n",
      "[126] time=0.12, avg_loss=1.0359, train_err=0.5231, test_h1=0.5526, test_l2=0.5566\n",
      "[129] time=0.12, avg_loss=1.0320, train_err=0.5212, test_h1=0.5524, test_l2=0.5528\n",
      "[132] time=0.12, avg_loss=1.0291, train_err=0.5197, test_h1=0.5522, test_l2=0.5503\n",
      "[135] time=0.12, avg_loss=1.0269, train_err=0.5186, test_h1=0.5514, test_l2=0.5475\n",
      "[138] time=0.12, avg_loss=1.0261, train_err=0.5182, test_h1=0.5501, test_l2=0.5460\n",
      "[141] time=0.12, avg_loss=1.0247, train_err=0.5175, test_h1=0.5504, test_l2=0.5478\n",
      "[144] time=0.12, avg_loss=1.0180, train_err=0.5141, test_h1=0.5533, test_l2=0.5346\n",
      "[147] time=0.12, avg_loss=1.0104, train_err=0.5103, test_h1=0.5574, test_l2=0.5530\n",
      "[150] time=0.12, avg_loss=1.0075, train_err=0.5088, test_h1=0.5555, test_l2=0.5366\n",
      "[153] time=0.12, avg_loss=1.0154, train_err=0.5128, test_h1=0.5545, test_l2=0.5415\n",
      "[156] time=0.12, avg_loss=1.0064, train_err=0.5082, test_h1=0.5554, test_l2=0.5430\n",
      "[159] time=0.12, avg_loss=1.0013, train_err=0.5056, test_h1=0.5541, test_l2=0.5357\n",
      "[162] time=0.12, avg_loss=0.9985, train_err=0.5043, test_h1=0.5523, test_l2=0.5368\n",
      "[165] time=0.12, avg_loss=0.9902, train_err=0.5001, test_h1=0.5523, test_l2=0.5370\n",
      "[168] time=0.12, avg_loss=0.9889, train_err=0.4994, test_h1=0.5531, test_l2=0.5310\n",
      "[171] time=0.12, avg_loss=0.9876, train_err=0.4987, test_h1=0.5538, test_l2=0.5359\n",
      "[174] time=0.12, avg_loss=0.9846, train_err=0.4972, test_h1=0.5566, test_l2=0.5314\n",
      "[177] time=0.12, avg_loss=0.9849, train_err=0.4974, test_h1=0.5574, test_l2=0.5334\n",
      "[180] time=0.12, avg_loss=0.9825, train_err=0.4961, test_h1=0.5598, test_l2=0.5360\n",
      "[183] time=0.12, avg_loss=0.9840, train_err=0.4969, test_h1=0.5604, test_l2=0.5323\n",
      "[186] time=0.12, avg_loss=0.9763, train_err=0.4931, test_h1=0.5690, test_l2=0.5365\n",
      "[189] time=0.12, avg_loss=0.9748, train_err=0.4923, test_h1=0.5724, test_l2=0.5393\n",
      "[192] time=0.12, avg_loss=0.9680, train_err=0.4888, test_h1=0.5707, test_l2=0.5412\n",
      "[195] time=0.12, avg_loss=0.9634, train_err=0.4865, test_h1=0.5715, test_l2=0.5390\n",
      "[198] time=0.12, avg_loss=0.9678, train_err=0.4887, test_h1=0.5715, test_l2=0.5500\n",
      "[201] time=0.12, avg_loss=0.9662, train_err=0.4879, test_h1=0.5791, test_l2=0.5483\n"
     ]
    }
   ],
   "source": [
    "tb_callback = TensorBoardCallback(log_dir=f'./logs/run{exp_num}')\n",
    "trainer = Trainer(model=model, n_epochs=202,\n",
    "                  device=device,\n",
    "                  wandb_log=False,\n",
    "                  log_test_interval=3,\n",
    "                  use_distributed=False,\n",
    "                  verbose=True,\n",
    "                  callbacks=[tb_callback])\n",
    "\n",
    "trainer.train(train_loader=train_loader, model=model, \n",
    "              output_encoder=None,\n",
    "              test_loaders=test_loader,\n",
    "              optimizer=optimizer,\n",
    "              scheduler=scheduler,\n",
    "              regularizer=False,\n",
    "              training_loss=train_loss,\n",
    "              eval_losses=eval_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Посмотрим на результаты экспериментов в Tensorboard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-12-02 13:33:06.382849: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-02 13:33:06.802559: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "2023-12-02 13:33:07.359897: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 13:33:07.360608: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 13:33:07.366117: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1960] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "\n",
      "NOTE: Using experimental fast data loading logic. To disable, pass\n",
      "    \"--load_fast=false\" and report issues on GitHub. More details:\n",
      "    https://github.com/tensorflow/tensorboard/issues/4784\n",
      "\n",
      "I1202 13:33:07.731471 139768739698432 plugin.py:429] Monitor runs begin\n",
      "Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\n",
      "TensorBoard 2.13.0 at http://localhost:6006/ (Press CTRL+C to quit)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!tensorboard --logdir=logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Графики для одного из запусков. Обратите внимание, функция потерь на обучении и тесте усреднена (reductions='mean'). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/1_train.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/1_test.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Выводы:**\n",
    "\n",
    "1. FNO достаточно хорошо справился с задачей, MSE на тесте составила около $4 \\cdot 10^{-2}$, что неплохо для данной задачи \\\n",
    "(т.к. нет никакой дополнительной физ.информации про начальные условия);\n",
    "\n",
    "2. Кривая потерь на обучении и тесте достаточно гладкая, что говорит о стабильном обучении FNO с подобранными гиперпараметрами."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Этап 2: написание собственной реализации FNO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "На данном этапе требуется реализовать FNO собственноручно и провести исследование о влиянии количества гармоник на качество обучения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.fno import MyFNO1d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проведем обучение с разным числом мод."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modes = [8, 16, 32, 48, 64, 96, 128]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for modes_num in modes:\n",
    "    print(f'Training custom FNO with {modes_num} modes')\n",
    "    print()\n",
    "    tb_callback = TensorBoardCallback(log_dir=f'./modes_research2/run_{modes_num}_modes')\n",
    "\n",
    "    my_model = MyFNO1d(n_modes=modes_num, in_channels=3, out_channels=1, n_layers=6, hidden_channels=32, projection_channels=64).to(device)\n",
    "    my_model = my_model.to(device)\n",
    "\n",
    "    optimizer = torch.optim.Adam(my_model.parameters(),\n",
    "                                lr=1e-3)\n",
    "\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10000, gamma=1)\n",
    "    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=5e-4)\n",
    "\n",
    "    h1loss = H1Loss(d=1, reductions='mean')\n",
    "    l2loss = LpLoss(d=1, p=2, reductions='mean')\n",
    "\n",
    "    train_loss = h1loss\n",
    "    eval_losses={'h1': h1loss, 'l2': l2loss}\n",
    "\n",
    "    trainer = Trainer(model=my_model, n_epochs=202,\n",
    "                  device=device,\n",
    "                  wandb_log=False,\n",
    "                  log_test_interval=3,\n",
    "                  use_distributed=False,\n",
    "                  verbose=True,\n",
    "                  callbacks=[tb_callback])\n",
    "\n",
    "    trainer.train(train_loader=train_loader, model=my_model, \n",
    "                output_encoder=None,\n",
    "                test_loaders=test_loader,\n",
    "                optimizer=optimizer,\n",
    "                scheduler=scheduler,\n",
    "                regularizer=False,\n",
    "                training_loss=train_loss,\n",
    "                eval_losses=eval_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Логи хранятся в директориях modes_research и modes_research2 (два запуска на одном наборе мод). На них можно взглянуть через tensorboard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-12-02 14:52:04.571522: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-02 14:52:04.990821: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "2023-12-02 14:52:05.539452: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 14:52:05.539913: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 14:52:05.544059: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1960] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "\n",
      "NOTE: Using experimental fast data loading logic. To disable, pass\n",
      "    \"--load_fast=false\" and report issues on GitHub. More details:\n",
      "    https://github.com/tensorflow/tensorboard/issues/4784\n",
      "\n",
      "I1202 14:52:05.910425 140576800544512 plugin.py:429] Monitor runs begin\n",
      "Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\n",
      "TensorBoard 2.13.0 at http://localhost:6006/ (Press CTRL+C to quit)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!tensorboard --logdir=modes_research"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Графики для одного из запусков:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/modes_train.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/modes_test_l2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/modes_test_h1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Выводы:**\n",
    "\n",
    "1. Собственная реализация FNO работает корректно и показывает такие же результаты, как FNO из neuralop;\n",
    "2. Кривая обучения на тесте у запусков с маленьким количеством мод (8, 16, 32) достаточно пологая и в итоге ошибка выше, чем у запусков с большим числом мод;\n",
    "3. На тесте сильно выделяется запуск с 128 модами (без фильтрации) - он дает наименьшую ошибку, результаты запусков с другим числом мод сопоставимы. \n",
    "    \n",
    "   Однако и здесь прослеживается та же закономерность - на запусках с большим числом мод лучше сходимость.\n",
    "\n",
    "**Итог:** в данных экспериментах число мод оказывало большое влияние на стабильность обучения и небольшое влияние на перфоманс модели (больше мод - лучше перфоманс и стабильнее обучение)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Этап 3 - сравнение с 1d-UNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для сравнения с FNO я выбрал архитектуру 1d-UNet."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/1dunet.png \"1D-UNet для размера входа 512\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Создадим модель и обучим её."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.unet import UNet1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "UNet has 549521 parameters.\n",
      "\n",
      "### MODEL ###\n",
      " UNet1d(\n",
      "  (encoder): ModuleList(\n",
      "    (0): EncoderBlock(\n",
      "      (conv1): Conv1d(3, 16, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm1): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv1d(16, 16, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (pooling): AvgPool1d(kernel_size=(3,), stride=(2,), padding=(1,))\n",
      "      (norm2): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (act): ReLU()\n",
      "    )\n",
      "    (1): EncoderBlock(\n",
      "      (conv1): Conv1d(16, 32, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm1): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv1d(32, 32, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (pooling): AvgPool1d(kernel_size=(3,), stride=(2,), padding=(1,))\n",
      "      (norm2): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (act): ReLU()\n",
      "    )\n",
      "    (2): EncoderBlock(\n",
      "      (conv1): Conv1d(32, 64, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (pooling): AvgPool1d(kernel_size=(3,), stride=(2,), padding=(1,))\n",
      "      (norm2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (act): ReLU()\n",
      "    )\n",
      "    (3): EncoderBlock(\n",
      "      (conv1): Conv1d(64, 128, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv1d(128, 128, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (pooling): AvgPool1d(kernel_size=(3,), stride=(2,), padding=(1,))\n",
      "      (norm2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (act): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (mid_layer): Sequential(\n",
      "    (0): Conv1d(128, 256, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "    (1): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "    (3): Conv1d(256, 256, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "    (4): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (5): ReLU()\n",
      "    (6): ConvTranspose1d(256, 128, kernel_size=(2,), stride=(2,))\n",
      "  )\n",
      "  (decoder): ModuleList(\n",
      "    (0): DecoderBlock(\n",
      "      (ch_reduce): Conv1d(256, 128, kernel_size=(1,), stride=(1,))\n",
      "      (conv): Conv1d(128, 64, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (upsample): ConvTranspose1d(64, 64, kernel_size=(2,), stride=(2,))\n",
      "      (act): ReLU()\n",
      "    )\n",
      "    (1): DecoderBlock(\n",
      "      (ch_reduce): Conv1d(128, 64, kernel_size=(1,), stride=(1,))\n",
      "      (conv): Conv1d(64, 32, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (upsample): ConvTranspose1d(32, 32, kernel_size=(2,), stride=(2,))\n",
      "      (act): ReLU()\n",
      "    )\n",
      "    (2): DecoderBlock(\n",
      "      (ch_reduce): Conv1d(64, 32, kernel_size=(1,), stride=(1,))\n",
      "      (conv): Conv1d(32, 16, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "      (norm): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (upsample): ConvTranspose1d(16, 16, kernel_size=(2,), stride=(2,))\n",
      "      (act): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (output_layer): Sequential(\n",
      "    (0): Conv1d(32, 16, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "    (1): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "    (3): Conv1d(16, 1, kernel_size=(1,), stride=(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: None\n",
      "    initial_lr: 0.0005\n",
      "    lr: 0.0005\n",
      "    maximize: False\n",
      "    weight_decay: 0\n",
      ")\n",
      "\n",
      "### SCHEDULER ###\n",
      " <torch.optim.lr_scheduler.StepLR object at 0x7f53ae921e50>\n",
      "\n",
      "### LOSSES ###\n",
      "\n",
      " * Train: <neuralop.training.losses.H1Loss object at 0x7f53ae9214c0>\n",
      "\n",
      " * Test: {'h1': <neuralop.training.losses.H1Loss object at 0x7f53ae9214c0>, 'l2': <neuralop.training.losses.LpLoss object at 0x7f53e56e2c10>}\n"
     ]
    }
   ],
   "source": [
    "st_ch=16\n",
    "unet = UNet1d(in_channels=3, start_channels=st_ch, out_channels=1)\n",
    "unet = unet.to(device)\n",
    "\n",
    "n_params = count_params(unet)\n",
    "print(f'\\nUNet has {n_params} parameters.')\n",
    "sys.stdout.flush()\n",
    "\n",
    "optimizer = torch.optim.Adam(unet.parameters(),\n",
    "                             lr=5e-4)\n",
    "\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)\n",
    "\n",
    "h1loss = H1Loss(d=1, reductions='mean')\n",
    "l2loss =LpLoss(d=1, p=2, reductions='mean')\n",
    "\n",
    "train_loss = h1loss\n",
    "eval_losses={'h1': h1loss, 'l2': l2loss}\n",
    "\n",
    "print('\\n### MODEL ###\\n', unet)\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",
    "sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "unet_exp_num = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using standard method to load data to device.\n",
      "using standard method to compute loss.\n",
      "Training on regular inputs (no multi-grid patching).\n",
      "Training on 400 samples\n",
      "Testing on [100] samples         on resolutions ['test'].\n",
      "Training on raw inputs of size x.shape=torch.Size([32, 3, 128]), y.shape=torch.Size([32, 1, 128])\n",
      ".. patched inputs of size x.shape=torch.Size([32, 3, 128]), y.shape=torch.Size([32, 1, 128])\n",
      "Raw outputs of size out.shape=torch.Size([32, 1, 128])\n",
      ".. Processed (unpatched) outputs of size out.shape=torch.Size([32, 1, 128])\n",
      "[0] time=0.14, avg_loss=0.1464, train_err=0.0739, test_h1=0.0701, test_l2=0.0744\n",
      "[3] time=0.13, avg_loss=0.0658, train_err=0.0332, test_h1=0.0716, test_l2=0.0679\n",
      "[6] time=0.13, avg_loss=0.0608, train_err=0.0307, test_h1=0.0636, test_l2=0.0625\n",
      "[9] time=0.13, avg_loss=0.0569, train_err=0.0287, test_h1=0.0580, test_l2=0.0612\n",
      "[12] time=0.13, avg_loss=0.0552, train_err=0.0279, test_h1=0.0578, test_l2=0.0620\n",
      "[15] time=0.13, avg_loss=0.0543, train_err=0.0274, test_h1=0.0575, test_l2=0.0622\n",
      "[18] time=0.13, avg_loss=0.0535, train_err=0.0270, test_h1=0.0570, test_l2=0.0618\n",
      "[21] time=0.13, avg_loss=0.0530, train_err=0.0268, test_h1=0.0568, test_l2=0.0613\n",
      "[24] time=0.13, avg_loss=0.0527, train_err=0.0266, test_h1=0.0568, test_l2=0.0607\n",
      "[27] time=0.13, avg_loss=0.0524, train_err=0.0265, test_h1=0.0571, test_l2=0.0596\n",
      "[30] time=0.13, avg_loss=0.0511, train_err=0.0258, test_h1=0.0592, test_l2=0.0614\n",
      "[33] time=0.13, avg_loss=0.0523, train_err=0.0264, test_h1=0.0595, test_l2=0.0585\n",
      "[36] time=0.13, avg_loss=0.0514, train_err=0.0260, test_h1=0.0559, test_l2=0.0571\n",
      "[39] time=0.13, avg_loss=0.0497, train_err=0.0251, test_h1=0.0542, test_l2=0.0571\n",
      "[42] time=0.13, avg_loss=0.0489, train_err=0.0247, test_h1=0.0546, test_l2=0.0561\n",
      "[45] time=0.12, avg_loss=0.0465, train_err=0.0235, test_h1=0.0554, test_l2=0.0551\n",
      "[48] time=0.13, avg_loss=0.0468, train_err=0.0236, test_h1=0.0554, test_l2=0.0561\n",
      "[51] time=0.12, avg_loss=0.0430, train_err=0.0217, test_h1=0.0503, test_l2=0.0538\n",
      "[54] time=0.13, avg_loss=0.0412, train_err=0.0208, test_h1=0.0486, test_l2=0.0510\n",
      "[57] time=0.13, avg_loss=0.0404, train_err=0.0204, test_h1=0.0517, test_l2=0.0510\n",
      "[60] time=0.13, avg_loss=0.0394, train_err=0.0199, test_h1=0.0493, test_l2=0.0492\n",
      "[63] time=0.13, avg_loss=0.0398, train_err=0.0201, test_h1=0.0512, test_l2=0.0530\n",
      "[66] time=0.13, avg_loss=0.0405, train_err=0.0205, test_h1=0.0541, test_l2=0.0530\n",
      "[69] time=0.13, avg_loss=0.0408, train_err=0.0206, test_h1=0.0501, test_l2=0.0488\n",
      "[72] time=0.13, avg_loss=0.0387, train_err=0.0195, test_h1=0.0494, test_l2=0.0500\n",
      "[75] time=0.13, avg_loss=0.0362, train_err=0.0183, test_h1=0.0479, test_l2=0.0471\n",
      "[78] time=0.13, avg_loss=0.0370, train_err=0.0187, test_h1=0.0510, test_l2=0.0491\n",
      "[81] time=0.13, avg_loss=0.0346, train_err=0.0175, test_h1=0.0476, test_l2=0.0497\n",
      "[84] time=0.13, avg_loss=0.0371, train_err=0.0187, test_h1=0.0495, test_l2=0.0474\n",
      "[87] time=0.13, avg_loss=0.0342, train_err=0.0173, test_h1=0.0511, test_l2=0.0498\n",
      "[90] time=0.13, avg_loss=0.0337, train_err=0.0170, test_h1=0.0467, test_l2=0.0467\n",
      "[93] time=0.13, avg_loss=0.0333, train_err=0.0168, test_h1=0.0492, test_l2=0.0476\n",
      "[96] time=0.13, avg_loss=0.0331, train_err=0.0167, test_h1=0.0513, test_l2=0.0482\n",
      "[99] time=0.13, avg_loss=0.0315, train_err=0.0159, test_h1=0.0489, test_l2=0.0477\n",
      "[102] time=0.13, avg_loss=0.0279, train_err=0.0141, test_h1=0.0474, test_l2=0.0467\n",
      "[105] time=0.13, avg_loss=0.0296, train_err=0.0150, test_h1=0.0469, test_l2=0.0447\n",
      "[108] time=0.13, avg_loss=0.0279, train_err=0.0141, test_h1=0.0503, test_l2=0.0453\n",
      "[111] time=0.13, avg_loss=0.0289, train_err=0.0146, test_h1=0.0481, test_l2=0.0458\n",
      "[114] time=0.13, avg_loss=0.0270, train_err=0.0136, test_h1=0.0495, test_l2=0.0460\n",
      "[117] time=0.13, avg_loss=0.0265, train_err=0.0134, test_h1=0.0481, test_l2=0.0458\n",
      "[120] time=0.12, avg_loss=0.0254, train_err=0.0128, test_h1=0.0508, test_l2=0.0473\n",
      "[123] time=0.14, avg_loss=0.0256, train_err=0.0129, test_h1=0.0508, test_l2=0.0459\n",
      "[126] time=0.13, avg_loss=0.0248, train_err=0.0125, test_h1=0.0483, test_l2=0.0451\n",
      "[129] time=0.13, avg_loss=0.0257, train_err=0.0130, test_h1=0.0500, test_l2=0.0457\n",
      "[132] time=0.13, avg_loss=0.0247, train_err=0.0125, test_h1=0.0492, test_l2=0.0459\n",
      "[135] time=0.13, avg_loss=0.0228, train_err=0.0115, test_h1=0.0489, test_l2=0.0442\n",
      "[138] time=0.13, avg_loss=0.0240, train_err=0.0121, test_h1=0.0484, test_l2=0.0438\n",
      "[141] time=0.13, avg_loss=0.0224, train_err=0.0113, test_h1=0.0476, test_l2=0.0447\n",
      "[144] time=0.13, avg_loss=0.0236, train_err=0.0119, test_h1=0.0479, test_l2=0.0453\n",
      "[147] time=0.13, avg_loss=0.0234, train_err=0.0118, test_h1=0.0482, test_l2=0.0440\n",
      "[150] time=0.13, avg_loss=0.0220, train_err=0.0111, test_h1=0.0490, test_l2=0.0452\n",
      "[153] time=0.13, avg_loss=0.0190, train_err=0.0096, test_h1=0.0482, test_l2=0.0442\n",
      "[156] time=0.13, avg_loss=0.0179, train_err=0.0091, test_h1=0.0485, test_l2=0.0439\n",
      "[159] time=0.13, avg_loss=0.0178, train_err=0.0090, test_h1=0.0487, test_l2=0.0438\n",
      "[162] time=0.13, avg_loss=0.0182, train_err=0.0092, test_h1=0.0487, test_l2=0.0439\n",
      "[165] time=0.13, avg_loss=0.0186, train_err=0.0094, test_h1=0.0488, test_l2=0.0443\n",
      "[168] time=0.13, avg_loss=0.0186, train_err=0.0094, test_h1=0.0488, test_l2=0.0445\n",
      "[171] time=0.13, avg_loss=0.0181, train_err=0.0091, test_h1=0.0485, test_l2=0.0439\n",
      "[174] time=0.13, avg_loss=0.0182, train_err=0.0092, test_h1=0.0491, test_l2=0.0441\n",
      "[177] time=0.13, avg_loss=0.0188, train_err=0.0095, test_h1=0.0495, test_l2=0.0450\n",
      "[180] time=0.13, avg_loss=0.0191, train_err=0.0096, test_h1=0.0490, test_l2=0.0443\n",
      "[183] time=0.13, avg_loss=0.0182, train_err=0.0092, test_h1=0.0491, test_l2=0.0444\n",
      "[186] time=0.13, avg_loss=0.0173, train_err=0.0087, test_h1=0.0488, test_l2=0.0445\n",
      "[189] time=0.13, avg_loss=0.0171, train_err=0.0086, test_h1=0.0483, test_l2=0.0432\n",
      "[192] time=0.13, avg_loss=0.0174, train_err=0.0088, test_h1=0.0486, test_l2=0.0438\n",
      "[195] time=0.13, avg_loss=0.0170, train_err=0.0086, test_h1=0.0484, test_l2=0.0434\n",
      "[198] time=0.13, avg_loss=0.0169, train_err=0.0086, test_h1=0.0487, test_l2=0.0438\n",
      "[201] time=0.13, avg_loss=0.0152, train_err=0.0077, test_h1=0.0484, test_l2=0.0435\n"
     ]
    }
   ],
   "source": [
    "tb_callback = TensorBoardCallback(log_dir=f'./unet_logs/run{unet_exp_num}_ch{st_ch}')\n",
    "\n",
    "trainer = Trainer(model=unet, n_epochs=202,\n",
    "                  device=device,\n",
    "                  wandb_log=False,\n",
    "                  log_test_interval=3,\n",
    "                  use_distributed=False,\n",
    "                  verbose=True,\n",
    "                  callbacks=[tb_callback])\n",
    "\n",
    "trainer.train(train_loader=train_loader, model=unet, \n",
    "            output_encoder=None,\n",
    "            test_loaders=test_loader,\n",
    "            optimizer=optimizer,\n",
    "            scheduler=scheduler,\n",
    "            regularizer=False,\n",
    "            training_loss=train_loss,\n",
    "            eval_losses=eval_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-12-02 14:43:19.679547: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-12-02 14:43:20.102595: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "2023-12-02 14:43:20.659169: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 14:43:20.659688: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-12-02 14:43:20.663977: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1960] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "\n",
      "NOTE: Using experimental fast data loading logic. To disable, pass\n",
      "    \"--load_fast=false\" and report issues on GitHub. More details:\n",
      "    https://github.com/tensorflow/tensorboard/issues/4784\n",
      "\n",
      "I1202 14:43:21.021307 140461344466688 plugin.py:429] Monitor runs begin\n",
      "Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\n",
      "TensorBoard 2.13.0 at http://localhost:6006/ (Press CTRL+C to quit)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!tensorboard --logdir=unet_logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ниже приведены кривые потерь на обучении и тесте для двух запусков с 64 и 32 входными каналами у архитектуры UNet. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/unet_train.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/unet_test_l2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/unet_test_h1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Выводы:**\n",
    "\n",
    "1. UNet обучается гораздо нестабильнее, чем FNO.\n",
    "2. В итоге UNet достигает примерно равного качества с FNO, однако в FNO меньше параметров (200-800K для 32-128 мод у FNO против 2.1M и 8.7M для 32 и 64 входных каналов у UNet), и качество FNO с 128 модами все еще лучше качества всех запусков UNet.\n",
    "\n",
    "**Итог**: в данной задаче FNO обучается более стабильно, чем UNet, и показывает лучшее качество при меньшем количестве параметров."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](imgs/fno_mem.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "task2",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}