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