{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from trnbl.training_manager import TrainingManager\n",
    "from trnbl.loggers.local import LocalLogger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "\tdef __init__(self) -> None:\n",
    "\t\tsuper(Model, self).__init__()\n",
    "\t\tself.fc: nn.Linear = nn.Linear(1, 1)\n",
    "\n",
    "\tdef forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "\t\treturn self.fc(x)\n",
    "\n",
    "\n",
    "class MockedDataset(torch.utils.data.Dataset):\n",
    "\tdef __init__(\n",
    "\t\tself,\n",
    "\t\tlength: int,\n",
    "\t\tchannels: int = 2,\n",
    "\t) -> None:\n",
    "\t\tself.dataset = torch.randn(length, channels, 1)\n",
    "\n",
    "\tdef __getitem__(self, idx: int):\n",
    "\t\treturn self.dataset[idx][0], self.dataset[idx][1]\n",
    "\n",
    "\tdef __len__(self):\n",
    "\t\treturn len(self.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# starting logger with id h40181-250116_2136-cenoza\n",
      "# starting training manager initialization\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "training run:   0%|          | 0/10 [00:00<?, ? epochs/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# initialized training manager\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1/10: 100%|██████████| 10/10 [00:00<00:00, 27.01 batches/s]\n",
      "epoch 2/10: 100%|██████████| 10/10 [00:00<00:00, 28.97 batches/s]\n",
      "epoch 3/10: 100%|██████████| 10/10 [00:00<00:00, 28.97 batches/s]\n",
      "epoch 4/10: 100%|██████████| 10/10 [00:00<00:00, 29.05 batches/s]\n",
      "epoch 5/10: 100%|██████████| 10/10 [00:00<00:00, 31.03 batches/s]\n",
      "epoch 6/10: 100%|██████████| 10/10 [00:00<00:00, 30.97 batches/s]\n",
      "epoch 7/10: 100%|██████████| 10/10 [00:00<00:00, 30.15 batches/s]\n",
      "epoch 8/10: 100%|██████████| 10/10 [00:00<00:00, 30.29 batches/s]\n",
      "epoch 9/10: 100%|██████████| 10/10 [00:00<00:00, 29.45 batches/s]\n",
      "epoch 10/10: 100%|██████████| 10/10 [00:00<00:00, 27.81 batches/s]\n",
      "training run: 100%|██████████| 10/10 [00:03<00:00,  2.89 epochs/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# training complete\n",
      "# closing logger\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model = Model()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.1)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "logger = LocalLogger(\n",
    "\tproject=\"integration-tests\",\n",
    "\tmetric_names=[\"train/loss\", \"train/acc\", \"val/loss\", \"val/acc\"],\n",
    "\ttrain_config=dict(\n",
    "\t\tmodel=str(model),\n",
    "\t\tdataset=\"dummy\",\n",
    "\t\toptimizer=str(optimizer),\n",
    "\t\tcriterion=str(criterion),\n",
    "\t),\n",
    "\tbase_path=\"tests/_temp\",\n",
    ")\n",
    "\n",
    "\n",
    "train_loader: DataLoader = DataLoader(MockedDataset(100), batch_size=10)\n",
    "\n",
    "with TrainingManager(\n",
    "\tmodel=model,\n",
    "\tlogger=logger,\n",
    "\tevals={\n",
    "\t\t\"1 epochs\": lambda model: {\"wgt_mean\": torch.mean(model.fc.weight).item()},\n",
    "\t\t\"1/2 epochs\": lambda model: logger.get_mem_usage(),\n",
    "\t}.items(),\n",
    "\tcheckpoint_interval=\"50 epochs\",\n",
    ") as tr:\n",
    "\t# Training loop\n",
    "\tfor epoch in tr.epoch_loop(range(10), use_tqdm=True):\n",
    "\t\tfor inputs, targets in tr.batch_loop(train_loader, use_tqdm=True):\n",
    "\t\t\toptimizer.zero_grad()\n",
    "\t\t\toutputs = model(inputs)\n",
    "\t\t\tloss = criterion(outputs, targets)\n",
    "\t\t\tloss.backward()\n",
    "\t\t\toptimizer.step()\n",
    "\n",
    "\t\t\taccuracy = torch.sum(torch.argmax(outputs, dim=1) == targets).item() / len(\n",
    "\t\t\t\ttargets\n",
    "\t\t\t)\n",
    "\n",
    "\t\t\ttr.batch_update(\n",
    "\t\t\t\tsamples=len(targets),\n",
    "\t\t\t\t**{\"train/loss\": loss.item(), \"train/acc\": accuracy},\n",
    "\t\t\t)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trnbl-XaCbE_mo-py3.12",
   "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.12.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
