Final fixes to get the sequence operational

Author: wannabeOG

encoder_train.py

@@ -164,6 +164,6 @@ def autoencoder_train(model, feature_extractor, path, optimizer, encoder_criteri
 
 	elapsed_time = time.time()-since
 	print ("This procedure took {:.2f} minutes and {:.2f} seconds".format(elapsed_time//60, elapsed_time%60))
-	print ("The best performing model has a {:.2f} loss on the test set".format(best_perform))
+	#print ("The best performing model has a {:.2f} loss on the test set".format(best_perform))
 	print ()
 

generate_models.py

@@ -39,7 +39,7 @@
 from model_train import *
 from model_utils import *
 
-from Initial_model_train import *
+from initial_model_train import *
 
 #define the parser
 parser = argparse.ArgumentParser(description='Generate models file')
@@ -167,6 +167,5 @@
 	#Train the model
 	if(task_number == 1):
 		train_model_1(len(image_folder.classes), feature_extractor, encoder_criterion, dset_loaders, dset_size, num_epochs_model , True, task_number,  lr = lr)
-	
 	else:	
 		train_model(len(image_folder.classes), feature_extractor, encoder_criterion, dset_loaders, dset_size, num_epochs_model , True, task_number,  lr = lr)

initial_model_train.py

@@ -1,3 +1,7 @@
+"""
+Module to train the model for the first task. Seperated from the rest of the code for the purpose of clarity
+"""
+
 #!/usr/bin/env python
 # coding: utf-8
 
@@ -13,7 +17,22 @@
 from model_utils import *
 
 def train_model_1(num_classes, feature_extractor, encoder_criterion, dset_loaders, dset_size, num_epochs, use_gpu, task_number, lr = 0.1, alpha = 0.01):
-
+	""" 
+	Inputs: 
+		1) num_classes = The number of classes in the new task  
+		2) feature_extractor = A reference to the feature extractor model  
+		3) encoder_criterion = The loss criterion for training the Autoencoder
+		4) dset_loaders = Dataset loaders for the model
+		5) dset_size = Size of the dataset loaders
+		6) num_of_epochs = Number of epochs for which the model needs to be trained
+		7) use_gpu = A flag which would be set if the user has a CUDA enabled device
+		8) task_number = A number which represents the task for which the model is being trained
+		9) lr = initial learning rate for the model
+		10) alpha = Tradeoff factor for the loss   
+
+	Function: Trains the model on the first task specifically
+		
+	"""
 	since = time.time()
 	best_perform = 10e6
 	device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -84,21 +103,11 @@ def train_model_1(num_classes, feature_extractor, encoder_criterion, dset_loader
 			optimizer.zero_grad()
 			#model_init.zero_grad()
 
-			# loss_1 only takes in the outputs from the nodes of the old classes 
-
-			#loss1_output = output[:, :num_of_classes_old]
-			#loss2_output = output[:, num_of_classes_old:]
-	
 			loss = model_criterion(output, labels, flag = "CE")
 
 			del labels
 			#del output
 
-			#total_loss = alpha*loss_1 + loss_2
-
-			#del loss_1
-			#del loss_2
-
 			loss.backward()
 			optimizer.step()
 
@@ -109,7 +118,7 @@ def train_model_1(num_classes, feature_extractor, encoder_criterion, dset_loader
 
 		print('Epoch Loss:{}'.format(epoch_loss))
 
-		if(epoch != 0 and epoch != num_of_epochs -1 and (epoch+1) % 10 == 0):
+		if(epoch != 0 and epoch != num_epochs -1 and (epoch+1) % 10 == 0):
 			epoch_file_name = os.path.join(mypath, str(epoch+1)+'.pth.tar')
 			torch.save({
 			'epoch': epoch,

model_train.py

@@ -151,6 +151,8 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 		ref_model.train(False)
 		ref_model.to(device)
 
+		#print (ref_model)
+
 		for param in model_init.Tmodel.classifier.parameters():
 			param.requires_grad = True
 
@@ -176,14 +178,18 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 		#Actually makes the changes to the model_init, so slightly redundant
 		print ("Initializing the model to be trained")
 		model_init = initialize_new_model(model_init, num_classes, num_of_classes_old)
-		model_init.to(device)
+		#print (model_init)
+		#model_init.to(device)
 		start_epoch = 0
 
 	#The training process format or LwF (Learning without Forgetting)
 	# Add the start epoch code 
 
 	if (best_relatedness > 0.85):
 
+		model_init.to(device)
+		ref_model.to(device)
+
 		print ("Using the LwF approach")
 		for epoch in range(start_epoch, num_epochs):			
 			since = time.time()
@@ -197,7 +203,7 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 
 			#scales the optimizer every 10 epochs 
 			optimizer = exp_lr_scheduler(optimizer, epoch, lr)
-			model_init = model_init.train(True)
+			#model_init = model_init.train(True)
 
 			for data in dset_loaders:
 				input_data, labels = data
@@ -212,32 +218,27 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 					input_data  = Variable(input_data)
 					labels = Variable(labels)
 
-				model_init.to(device)
-				ref_model.to(device)
-				
 				output = model_init(input_data)
 				ref_output = ref_model(input_data)
-
 				del input_data
 
 				optimizer.zero_grad()
-				model_init.zero_grad()
 
 				# loss_1 only takes in the outputs from the nodes of the old classes 
 
 				loss1_output = output[:, :num_of_classes_old]
 				loss2_output = output[:, num_of_classes_old:]
 
+				print ()
+
 				del output
 
 				loss_1 = model_criterion(loss1_output, ref_output, flag = "Distill")
-				
 				del ref_output
 
 				# loss_2 takes in the outputs from the nodes that were initialized for the new task
 
 				loss_2 = model_criterion(loss2_output, labels, flag = "CE")
-				
 				del labels
 				#del output
 
@@ -257,7 +258,7 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 
 			print('Epoch Loss:{}'.format(epoch_loss))
 
-			if(epoch != 0 and epoch != num_of_epochs -1 and (epoch+1) % 10 == 0):
+			if(epoch != 0 and epoch != num_epochs -1 and (epoch+1) % 10 == 0):
 				epoch_file_name = os.path.join(mypath, str(epoch+1)+'.pth.tar')
 				torch.save({
 				'epoch': epoch,
@@ -277,6 +278,7 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 	#Process for finetuning the model
 	else:
 
+		model_init.to(device)
 		print ("Using the finetuning approach")
 
 		for epoch in range(start_epoch, num_epochs):
@@ -302,9 +304,6 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 					input_data  = Variable(input_data)
 					labels = Variable(labels)
 
-				#Shifts the model to the device
-				model_init.to(device)
-
 				output = model_init(input_data)
 
 				del input_data
@@ -314,7 +313,7 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 				model_init.zero_grad()
 
 				#Implemented as explained in the doc string
-				loss = model_criterion(output[num_of_classes_old:], labels)
+				loss = model_criterion(output[num_of_classes_old:], labels, flag = 'CE')
 
 				del output
 				del labels
@@ -330,7 +329,7 @@ def train_model(num_classes, feature_extractor, encoder_criterion, dset_loaders,
 
 			print('Epoch Loss:{}'.format(epoch_loss))
 
-			if(epoch != 0 and (epoch+1) % 5 == 0 and epoch != num_of_epochs -1):
+			if(epoch != 0 and (epoch+1) % 5 == 0 and epoch != num_epochs -1):
 				epoch_file_name = os.path.join(path_to_model, str(epoch+1)+'.pth.tar')
 				torch.save({
 				'epoch': epoch,