{ "cells": [ { "cell_type": "markdown", "id": "991cf177", "metadata": {}, "source": [ "# Training the Channel Estimation Model\n", "\n", "This notebook trains the deep neural network used for multi-layer channel estimation. It loads the training, validation, and test datasets generated in the previous step, initializes the `ChEstNet` model, and optimizes the network using mean squared error (MSE) loss.\n", "\n", "During training, the learning rate decays exponentially from the initial value to the final value over the specified number of epochs. After each epoch, the model is evaluated on the validation dataset. The checkpoint with the lowest validation loss is saved to the `Models` directory and reloaded at the end of training.\n", "\n", "An already-trained model is included in the `Models` directory. You may either run this notebook to train a new model or proceed directly to the evaluation notebooks using the provided model." ] }, { "cell_type": "code", "execution_count": 1, "id": "c90d6619", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import os\n", "import torch\n", "from torch.optim.lr_scheduler import ExponentialLR\n", "\n", "from ChEstNet import ChEstNet, ChEstDataset" ] }, { "cell_type": "code", "execution_count": 2, "id": "3bfa5fe8-e489-4612-bc5c-5c775a132c02", "metadata": {}, "outputs": [], "source": [ "# Load the datasets:\n", "dataPath = \"/data/datasets/SelfRefine\" # Replace with the path to your dataset files\n", "batchSize = 64 \n", "trainDS = ChEstDataset( os.path.join(dataPath,\"Train.npy\"), batchSize )\n", "validDS = ChEstDataset( os.path.join(dataPath,\"Valid.npy\"), batchSize )\n", "testDS = ChEstDataset( os.path.join(dataPath,\"Test.npy\"), batchSize )\n", " " ] }, { "cell_type": "code", "execution_count": 3, "id": "258684bd-7f88-4252-8066-d6dc7bb58ce2", "metadata": {}, "outputs": [], "source": [ "modelFileName = \"Models/Trained.pth\" # Output filename for the trained model\n", "numEpochs = 100\n", "lrStart, lrEnd = 0.002, 0.00002 # Learning rate decays exponentially from 'lrStart' to 'lrEnd'\n", "device = f\"cuda:0\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n", "model = ChEstNet(device) # Create the model\n", "\n", "optimizer = torch.optim.Adam(model.parameters(), lr=lrStart)\n", "lrScheduler = ExponentialLR(optimizer, np.exp(np.log(lrEnd/lrStart)/(numEpochs-1)))\n", "lossFunction = torch.nn.MSELoss()\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "b5861520-5e09-4691-b03b-a911a8531519", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Epoch Learning Rate Training Loss Validation Loss\n", "----- ------------- ------------- ---------------\n", " 1 0.002000 0.593710 0.388976 \n", " 2 0.001909 0.276281 0.240636 * \n", " 3 0.001822 0.225470 0.194113 * \n", " 4 0.001739 0.191013 0.177125 * \n", " 5 0.001660 0.174355 0.196268 \n", " 6 0.001585 0.162672 0.139205 * \n", " 7 0.001513 0.151027 0.153683 \n", " 8 0.001444 0.144716 0.140416 \n", " 9 0.001379 0.137015 0.119776 * \n", " 10 0.001316 0.130708 0.102059 * \n", " 11 0.001256 0.127196 0.107508 \n", " 12 0.001199 0.122850 0.111060 \n", " 13 0.001144 0.121160 0.110864 \n", " 14 0.001092 0.117182 0.106913 \n", " 15 0.001043 0.115687 0.098550 * \n", " 16 0.000995 0.113761 0.102031 \n", " 17 0.000950 0.111562 0.103003 \n", " 18 0.000907 0.110321 0.100105 \n", " 19 0.000866 0.109446 0.101697 \n", " 20 0.000826 0.107634 0.097564 * \n", " 21 0.000789 0.105831 0.095171 * \n", " 22 0.000753 0.104648 0.094937 * \n", " 23 0.000719 0.103204 0.096389 \n", " 24 0.000686 0.102920 0.092838 * \n", " 25 0.000655 0.101974 0.091227 * \n", " 26 0.000625 0.101095 0.091867 \n", " 27 0.000597 0.100181 0.089986 * \n", " 28 0.000570 0.099703 0.088168 * \n", " 29 0.000544 0.098449 0.088144 * \n", " 30 0.000519 0.097884 0.087780 * \n", " 31 0.000495 0.097204 0.086695 * \n", " 32 0.000473 0.096963 0.086214 * \n", " 33 0.000451 0.096408 0.090265 \n", " 34 0.000431 0.095796 0.086883 \n", " 35 0.000411 0.095553 0.087128 \n", " 36 0.000393 0.095162 0.088318 \n", " 37 0.000375 0.094596 0.086862 \n", " 38 0.000358 0.094252 0.086169 * \n", " 39 0.000341 0.093961 0.084881 * \n", " 40 0.000326 0.093534 0.085342 \n", " 41 0.000311 0.093005 0.084186 * \n", " 42 0.000297 0.092960 0.085097 \n", " 43 0.000283 0.092543 0.083574 * \n", " 44 0.000271 0.092341 0.084721 \n", " 45 0.000258 0.092323 0.084242 \n", " 46 0.000247 0.091788 0.083811 \n", " 47 0.000235 0.091534 0.084309 \n", " 48 0.000225 0.091448 0.083943 \n", " 49 0.000214 0.091001 0.083925 \n", " 50 0.000205 0.090791 0.083822 \n", " 51 0.000195 0.090719 0.082809 * \n", " 52 0.000187 0.090252 0.083022 \n", " 53 0.000178 0.090632 0.083907 \n", " 54 0.000170 0.090226 0.083190 \n", " 55 0.000162 0.089889 0.082961 \n", " 56 0.000155 0.089700 0.082639 * \n", " 57 0.000148 0.089480 0.082727 \n", " 58 0.000141 0.089512 0.082371 * \n", " 59 0.000135 0.089439 0.082200 * \n", " 60 0.000129 0.089024 0.082607 \n", " 61 0.000123 0.089087 0.082157 * \n", " 62 0.000117 0.089033 0.081656 * \n", " 63 0.000112 0.088912 0.081966 \n", " 64 0.000107 0.088967 0.081855 \n", " 65 0.000102 0.088824 0.081713 \n", " 66 0.000097 0.088641 0.081922 \n", " 67 0.000093 0.088577 0.081618 * \n", " 68 0.000089 0.088507 0.082051 \n", " 70 0.000081 0.088252 0.081858 \n", " 71 0.000077 0.088167 0.081529 * \n", " 72 0.000074 0.088039 0.081552 \n", " 73 0.000070 0.088081 0.081714 \n", " 74 0.000067 0.088052 0.081313 * \n", " 75 0.000064 0.087745 0.081495 \n", " 76 0.000061 0.087764 0.081391 \n", " 77 0.000058 0.087846 0.081122 * \n", " 78 0.000056 0.087613 0.081238 \n", " 79 0.000053 0.087572 0.081253 \n", " 80 0.000051 0.087519 0.081218 \n", " 81 0.000048 0.087531 0.081289 \n", " 82 0.000046 0.087567 0.081146 \n", " 83 0.000044 0.087631 0.081352 \n", " 84 0.000042 0.087436 0.081125 \n", " 85 0.000040 0.087425 0.081201 \n", " 86 0.000038 0.087305 0.081191 \n", " 87 0.000037 0.087241 0.081107 * \n", " 88 0.000035 0.087375 0.081080 * \n", " 89 0.000033 0.087215 0.081038 * \n", " 90 0.000032 0.087309 0.081036 * \n", " 91 0.000030 0.087294 0.081164 \n", " 92 0.000029 0.087123 0.081232 \n", " 93 0.000028 0.087042 0.080875 * \n", " 94 0.000026 0.087075 0.081113 \n", " 95 0.000025 0.086979 0.080928 \n", " 96 0.000024 0.086935 0.080861 * \n", " 97 0.000023 0.087059 0.080926 \n", " 98 0.000022 0.086853 0.080815 * \n", " 99 0.000021 0.086966 0.080932 \n", " 100 0.000020 0.086941 0.080734 * \n" ] } ], "source": [ "# Main training loop:\n", "lowestLoss, bestEpoch = None, None\n", "validLoss = None\n", "print(\"Epoch Learning Rate Training Loss Validation Loss\")\n", "print(\"----- ------------- ------------- ---------------\")\n", "for epoch in range(numEpochs):\n", " curLr = lrScheduler.get_last_lr()[0]\n", " print(f\" {epoch+1:-4d} {curLr:-10f} \", end=\"\")\n", " \n", " # Train one epoch\n", " lossMin, lossMean, lossMax = model.trainEpoch(trainDS, lossFunction, optimizer)\n", " print(f\"{lossMean:-10f} \", end=\"\")\n", " \n", " validLoss = model.evaluate(validDS, lossFunction)\n", " if lowestLoss is None:\n", " lowestLoss, bestEpoch = validLoss, epoch+1\n", " model.saveParams(modelFileName) # Save the best model so far \n", " print(f\"{validLoss:-10f} \")\n", " elif validLoss