{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ca08df84",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0/1000, Loss: 1.0013\n",
      "Step 10/1000, Loss: 1.0088\n",
      "Step 20/1000, Loss: 0.9956\n",
      "Step 30/1000, Loss: 0.9781\n",
      "Step 40/1000, Loss: 0.9613\n",
      "Step 50/1000, Loss: 0.9313\n",
      "Step 60/1000, Loss: 0.8927\n",
      "Step 70/1000, Loss: 0.8503\n",
      "Step 80/1000, Loss: 0.7537\n",
      "Step 90/1000, Loss: 0.6689\n",
      "Step 100/1000, Loss: 0.6063\n",
      "Step 110/1000, Loss: 0.5172\n",
      "Step 120/1000, Loss: 0.4592\n",
      "Step 130/1000, Loss: 0.4044\n",
      "Step 140/1000, Loss: 0.3610\n",
      "Step 150/1000, Loss: 0.3175\n",
      "Step 160/1000, Loss: 0.2825\n",
      "Step 170/1000, Loss: 0.2560\n",
      "Step 180/1000, Loss: 0.2360\n",
      "Step 190/1000, Loss: 0.2203\n",
      "Step 200/1000, Loss: 0.1930\n",
      "Step 210/1000, Loss: 0.1854\n",
      "Step 220/1000, Loss: 0.1723\n",
      "Step 230/1000, Loss: 0.1546\n",
      "Step 240/1000, Loss: 0.1386\n",
      "Step 250/1000, Loss: 0.1271\n",
      "Step 260/1000, Loss: 0.1109\n",
      "Step 270/1000, Loss: 0.1032\n",
      "Step 280/1000, Loss: 0.0899\n",
      "Step 290/1000, Loss: 0.0807\n",
      "Step 300/1000, Loss: 0.0750\n",
      "Step 310/1000, Loss: 0.0813\n",
      "Step 320/1000, Loss: 0.0612\n",
      "Step 330/1000, Loss: 0.0544\n",
      "Step 340/1000, Loss: 0.0552\n",
      "Step 350/1000, Loss: 0.0446\n",
      "Step 360/1000, Loss: 0.0403\n",
      "Step 370/1000, Loss: 0.0350\n",
      "Step 380/1000, Loss: 0.0612\n",
      "Step 390/1000, Loss: 0.0364\n",
      "Step 400/1000, Loss: 0.0322\n",
      "Step 410/1000, Loss: 0.0302\n",
      "Step 420/1000, Loss: 0.0519\n",
      "Step 430/1000, Loss: 0.0319\n",
      "Step 440/1000, Loss: 0.0260\n",
      "Step 450/1000, Loss: 0.0208\n",
      "Step 460/1000, Loss: 0.0409\n",
      "Step 470/1000, Loss: 0.0291\n",
      "Step 480/1000, Loss: 0.0234\n",
      "Step 490/1000, Loss: 0.0194\n",
      "Step 500/1000, Loss: 0.0274\n",
      "Step 510/1000, Loss: 0.0231\n",
      "Step 520/1000, Loss: 0.0199\n",
      "Step 530/1000, Loss: 0.0154\n",
      "Step 540/1000, Loss: 0.0278\n",
      "Step 550/1000, Loss: 0.0185\n",
      "Step 560/1000, Loss: 0.0180\n",
      "Step 570/1000, Loss: 0.0152\n",
      "Step 580/1000, Loss: 0.0132\n",
      "Step 590/1000, Loss: 0.0111\n",
      "Step 600/1000, Loss: 0.0396\n",
      "Step 610/1000, Loss: 0.0179\n",
      "Step 620/1000, Loss: 0.0148\n",
      "Step 630/1000, Loss: 0.0123\n",
      "Step 640/1000, Loss: 0.0265\n",
      "Step 650/1000, Loss: 0.0133\n",
      "Step 660/1000, Loss: 0.0128\n",
      "Step 670/1000, Loss: 0.0107\n",
      "Step 680/1000, Loss: 0.0142\n",
      "Step 690/1000, Loss: 0.0202\n",
      "Step 700/1000, Loss: 0.0125\n",
      "Step 710/1000, Loss: 0.0107\n",
      "Step 720/1000, Loss: 0.0140\n",
      "Step 730/1000, Loss: 0.0195\n",
      "Step 740/1000, Loss: 0.0148\n",
      "Step 750/1000, Loss: 0.0109\n",
      "Step 760/1000, Loss: 0.0094\n",
      "Step 770/1000, Loss: 0.0121\n",
      "Step 780/1000, Loss: 0.0233\n",
      "Step 790/1000, Loss: 0.0151\n",
      "Step 800/1000, Loss: 0.0134\n",
      "Step 810/1000, Loss: 0.0117\n",
      "Step 820/1000, Loss: 0.0124\n",
      "Step 830/1000, Loss: 0.0221\n",
      "Step 840/1000, Loss: 0.0161\n",
      "Step 850/1000, Loss: 0.0136\n",
      "Step 860/1000, Loss: 0.0161\n",
      "Step 870/1000, Loss: 0.0194\n",
      "Step 880/1000, Loss: 0.0145\n",
      "Step 890/1000, Loss: 0.0149\n",
      "Step 900/1000, Loss: 0.0232\n",
      "Step 910/1000, Loss: 0.0166\n",
      "Step 920/1000, Loss: 0.0156\n",
      "Step 930/1000, Loss: 0.0276\n",
      "Step 940/1000, Loss: 0.0176\n",
      "Step 950/1000, Loss: 0.0152\n",
      "Step 960/1000, Loss: 0.0162\n",
      "Step 970/1000, Loss: 0.0143\n",
      "Step 980/1000, Loss: 0.0136\n",
      "Step 990/1000, Loss: 0.0117\n",
      "Total time: 67.25 s\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import time\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using device:\", device)\n",
    "\n",
    "\n",
    "# Огромные параметры\n",
    "N, D_in, H1, H2, H3, D_out = 300_000, 4096, 2048, 1024, 512, 10\n",
    "batch_size = 16_384  # большой батч\n",
    "steps = 1000          # много итераций для длительной нагрузки\n",
    "\n",
    "# Случайные данные на GPU\n",
    "x = torch.randn(N, D_in, device=device, dtype=torch.float32)\n",
    "y = torch.randn(N, D_out, device=device, dtype=torch.float32)\n",
    "\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(D_in, H1),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(H1, H2),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(H2, H3),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(H3, D_out)\n",
    ").to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "start = time.time()\n",
    "for t in range(steps):\n",
    "    idx = torch.randint(0, N, (batch_size,), device=device)\n",
    "    x_batch = x[idx]\n",
    "    y_batch = y[idx]\n",
    "\n",
    "    y_pred = model(x_batch)\n",
    "    loss = loss_fn(y_pred, y_batch)\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if t % 10 == 0:\n",
    "        # замедляем вывод, чтобы можно было наблюдать\n",
    "        print(f\"Step {t}/{steps}, Loss: {loss.item():.4f}\")\n",
    "\n",
    "end = time.time()\n",
    "print(f\"Total time: {end-start:.2f} s\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}