Here’s the full code to reproduce the issue.
Model and function definitions:
import matplotlib.pyplot as plt
class Conv1DAE(torch.nn.Module):
def __init__(self):
super().__init__()
self.Conv1D_L1 = torch.nn.Conv1d(7, 32, 48, 1)
self.LeakyReLU_L2 = torch.nn.LeakyReLU()
self.BatchNorm1D_L3 = torch.nn.BatchNorm1d(32)
self.MaxPool1D_L4 = torch.nn.MaxPool1d(4, 4, 2, return_indices = True)
self.MaxUnPool1D_L5 = torch.nn.MaxUnpool1d(4, 4)
self.Conv1DTransposed_L6 = torch.nn.ConvTranspose1d(32, 7, 45, 1)
def forward(self, InputData):
ConvTensor = self.Conv1D_L1(InputData)
ReLUTensor = self.LeakyReLU_L2(ConvTensor)
BatchNormTensor = self.BatchNorm1D_L3(ReLUTensor)
MaxPooledTensor, Indices = self.MaxPool1D_L4(BatchNormTensor)
MaxUnPooledTensor = self.MaxUnPool1D_L5(MaxPooledTensor, Indices)
Output = self.Conv1DTransposed_L6(MaxUnPooledTensor)
return Output
class ConvNet(torch.nn.Module):
def __init__(self):
super().__init__()
self.Conv1D_L1 = torch.nn.Conv1d(7, 32, 48, 1)
self.LeakyReLU_L2 = torch.nn.LeakyReLU()
self.BatchNorm1D_L3 = torch.nn.BatchNorm1d(32)
self.MaxPool1D_L4 = torch.nn.MaxPool1d(4, 4, 2)
self.Flattening_L5 = torch.nn.Flatten()
self.Linear_L6 = torch.nn.Linear(800, 400)
self.LeakyReLu_L7 = torch.nn.LeakyReLU()
self.DropOut_L8 = torch.nn.Dropout(0.2)
self.FinalLayer = torch.nn.Linear(400, 5)
def forward(self, InputData):
ConvHeadTensor = self.Conv1D_L1(InputData)
ReLUTensor1 = self.LeakyReLU_L2(ConvHeadTensor)
BatchNormTensor = self.BatchNorm1D_L3(ReLUTensor1)
MaxPoolTensor = self.MaxPool1D_L4(BatchNormTensor)
FlattenedTensor = self.Flattening_L5(MaxPoolTensor)
LinearTensor1 = self.Linear_L6(FlattenedTensor)
ReLUTensor2 = self.LeakyReLu_L7(LinearTensor1)
DropOutTensor = self.DropOut_L8(ReLUTensor2)
Output = self.FinalLayer(DropOutTensor)
return Output
def TrainAE(TrainLoader, TestLoader, AutoEncoder, NumberOfEpochs, LearningRate = 0.001):
LossMetric = torch.nn.MSELoss()
Optimiser = torch.optim.Adam(AutoEncoder.parameters(), lr = LearningRate)
if(torch.cuda.is_available()):
AutoEncoder.to("cuda:0")
UseGPU = True
else:
UseGPU = False
TrainingLoss = []
ValidationLoss = []
for Epoch in range(NumberOfEpochs):
TrainingRunningLoss = 0
AutoEncoder.train()
for Iteration, Data in enumerate(TrainLoader):
TrainX, _ = Data
if(UseGPU == True):
TrainX = TrainX.to(Device)
Optimiser.zero_grad()
Output = AutoEncoder(TrainX)
Loss = LossMetric(Output, TrainX)
TrainingRunningLoss += Loss.item()
Loss.backward()
Optimiser.step()
TrainingLoss.append(TrainingRunningLoss)
print("Epoch: {0}/{1}, Training Loss: {2}".format((Epoch + 1), NumberOfEpochs, round(TrainingRunningLoss, 5)))
ValidationRunningLoss = 0
AutoEncoder.eval()
with torch.no_grad():
for Iteration, Data in enumerate(TestLoader):
TestX, _ = Data
if(UseGPU == True):
TestX = TestX.to(Device)
Predictions = AutoEncoder(TestX)
Loss = LossMetric(Predictions, TestX)
ValidationRunningLoss += Loss.item()
ValidationLoss.append(ValidationRunningLoss)
print("Epoch: {0}/{1}, Validation Loss: {2}".format((Epoch + 1), NumberOfEpochs, round(ValidationRunningLoss, 5)))
plt.title("Plot of Auto Encoder losses")
plt.plot(TrainingLoss)
plt.plot(ValidationLoss)
plt.show()
plt.clf()
AutoEncoder.to("cpu")
AutoEncoderParameters = AutoEncoder.state_dict()
del(AutoEncoderParameters["Conv1DTransposed_L6.weight"])
del(AutoEncoderParameters["Conv1DTransposed_L6.bias"])
return AutoEncoderParameters
def TrainModel(TrainLoader, TestLoader, Model, NumberOfEpochs, Optimiser, LossMetric):
if(torch.cuda.is_available()):
Model.to("cuda:0")
UseGPU = True
else:
UseGPU = False
TrainingLoss = []
ValidationLoss = []
for Epoch in range(NumberOfEpochs):
TrainingRunningLoss = 0
Model.train()
for Iteration, Data in enumerate(TrainLoader):
TrainX, TrainY = Data
if(UseGPU == True):
TrainX = TrainX.to(Device)
TrainY = TrainY.to(Device)
Optimiser.zero_grad()
Output = Model(TrainX)
Loss = LossMetric(Output, TrainY)
TrainingRunningLoss += Loss.item()
Loss.backward()
Optimiser.step()
TrainingLoss.append(TrainingRunningLoss)
print("Epoch: {0}/{1}, Training Loss: {2}".format((Epoch + 1), NumberOfEpochs, TrainingRunningLoss))
ValidationRunningLoss = 0
Model.eval()
with torch.no_grad():
for Iteration, Data in enumerate(TestLoader):
TestX, TestY = Data
if(UseGPU == True):
TestX = TestX.to(Device)
TestY = TestY.to(Device)
Predictions = Model(TestX)
Loss = LossMetric(Predictions, TestY)
ValidationRunningLoss += Loss.item()
ValidationLoss.append(ValidationRunningLoss)
print("Epoch: {0}/{1}, Validation Loss: {2}".format((Epoch + 1), NumberOfEpochs, ValidationRunningLoss))
plt.title("Plot of training and validation losses")
plt.plot(TrainingLoss)
plt.plot(ValidationLoss)
plt.show()
plt.clf()
Model.to("cpu")
return Model
Weight loading along with print statements for visual checks:
AutoEncoder = Conv1DAE()
AutoEncoderParameters = TrainAE(TrainLoader, TestLoader, AutoEncoder, 325, LearningRate = 0.0000064)
MainNetwork = ConvNet()
MainNetworkParameters = MainNetwork.state_dict()
print(MainNetworkParameters["BatchNorm1D_L3.bias"])
print(AutoEncoderParameters["BatchNorm1D_L3.bias"])
MainNetworkParameters.update(AutoEncoderParameters)
MainNetwork.load_state_dict(MainNetworkParameters)
UpdatedMainNetworkParameters = MainNetwork.state_dict()
print(UpdatedMainNetworkParameters["BatchNorm1D_L3.bias"])
Setting the requires_grad
attribute of the pretrained layers to False
MainNetwork.Conv1D_L1.requires_grad = False
MainNetwork.BatchNorm1D_L3.requires_grad = False
Training of the main network:
LossMetric = torch.nn.CrossEntropyLoss()
Optimiser = torch.optim.Adam(MainNetwork.parameters(), lr = 0.0000001)
MainNetwork = TrainModel(TrainLoader, TestLoader, MainNetwork, 350, Optimiser, LossMetric)
Checking of the main network’s weights post training:
TrainedMainNetworkParameters = MainNetwork.state_dict()
print(TrainedMainNetworkParameters["BatchNorm1D_L3.bias"])