System and method of training a student model using a teacher model

Abstract

The disclosure relates to a method and system of training a student model using a teacher model. The method includes receiving, from a user, a selection of a target knowledge distillation technique, a teacher model, a student model, and one or more batches of training data. The method further includes loading the teacher model on a memory device and extracting knowledge output from the teacher model for each of the one or more batches of the training data, based on the target knowledge distillation technique, and sequentially storing extracted knowledge output in a knowledge database. The method further includes unloading the teacher model and loading the student model on the memory device, and training the student model based on ground-truth labels associated with each of the one or more batches of training data and the knowledge output corresponding to the target knowledge distillation technique.

What is claimed is: 
     
         1 . A method of training a student model using a teacher model, the method comprising:
 receiving, from a user, a selection of a target knowledge distillation technique from a plurality of knowledge distillation techniques, a teacher model, a student model, and one or more batches of training data;   loading the teacher model on a memory device;   extracting knowledge output from the teacher model for each of the one or more batches of the training data, based on the target knowledge distillation technique, and sequentially storing extracted knowledge output in a knowledge database, wherein the teacher model is a pre-trained model;   Upon extracting, unloading the teacher model from the memory device and loading the student model on the memory device; and   training the student model based on ground-truth labels associated with each of the one or more batches of training data and the knowledge output corresponding to the target knowledge distillation technique, by fetching the knowledge output corresponding to the target knowledge distillation technique and each of the one or more batches of training data from the knowledge database.   
     
     
         2 . The method of  claim 1 , wherein the plurality of knowledge distillation techniques comprises: a response-based knowledge distillation technique, a feature-based knowledge distillation technique, and a relation-based knowledge distillation technique. 
     
     
         3 . The method of  claim 2 , wherein the knowledge output corresponding to the response-based knowledge distillation technique comprises soft targets obtained from a final output layer of the teacher model. 
     
     
         4 . The method of  claim 3  further comprising:
 calculating distillation loss for the response-based knowledge distillation technique, using at least one of: a cross-entropy loss on the ground-truth labels associated with the training data and Kullback-Leibler (KL) divergence between the predictions from the teacher model and the predictions from the student model; and 
 adjusting weights of the student model, based on the distillation loss for the response-based knowledge distillation technique. 
 
     
     
         5 . The method of  claim 2 , wherein the knowledge output corresponding to the feature-based knowledge distillation technique comprises internal feature representations obtained from one or more intermediate layers of the teacher model. 
     
     
         6 . The method of  claim 5  further comprising:
 calculating distillation loss for the feature-based knowledge distillation technique, using at least one of: a Euclidean distance or cosine similarity between features of the teacher model and the student model, a mean squared error (MSE) loss, or a correlation alignment loss; and 
 adjusting weights of the student model, based on the distillation loss for feature-based knowledge distillation technique. 
 
     
     
         7 . The method of  claim 2 , wherein the knowledge output corresponding to the relation-based knowledge distillation technique comprises pair-wise relations and group-wise relations between data points obtained from the teacher model. 
     
     
         8 . The method of  claim 2  further comprising:
 calculating distillation loss for the relation-based knowledge distillation technique, by minimizing discrepancy between class relationships learned by the teacher model and the student model respectively; and 
 adjusting weights of the student model, based on the distillation loss for feature-based knowledge distillation technique. 
 
     
     
         9 . The method of  claim 1 , wherein training the student model comprises: iteratively inputting each of the one or more batches of training data to the student model based on a predefined epoch. 
     
     
         10 . A system for training a student model using a teacher model, the system comprising:
 a processor; and   a memory communicatively coupled to the processor, the memory storing a plurality of processor-executable instructions, wherein the processor-executable instructions, upon execution by the processor, cause the processor to:
 receive, from a user, a selection of a target knowledge distillation technique from a plurality of knowledge distillation techniques, a teacher model, a student model, and one or more batches of training data; 
 load the teacher model on a memory device; 
 extract knowledge output from the teacher model for each of the one or more batches of the training data, based on the target knowledge distillation technique, and sequentially storing extracted knowledge output in a knowledge database, wherein the teacher model is a pre-trained model; 
 upon extracting, unload the teacher model from the memory device and loading the student model on the memory device; and 
 train the student model based on ground-truth labels associated with each of the one or more batches of training data and the knowledge output corresponding to the target knowledge distillation technique, by fetching the knowledge output corresponding to the target knowledge distillation technique and each of the one or more batches of training data from the knowledge database. 
   
     
     
         11 . The system of  claim 10 , wherein the plurality of knowledge distillation techniques comprises: a response-based knowledge distillation technique, a feature-based knowledge distillation technique, and a relation-based knowledge distillation technique. 
     
     
         12 . The system of  claim 11 , wherein the knowledge output corresponding to the response-based knowledge distillation technique comprise soft targets obtained from a final output layer of the teacher model, and wherein the processor-executable instructions further cause the processor to:
 calculate distillation loss for the response-based knowledge distillation technique, using at least one of: a cross-entropy loss on the ground-truth labels associated with the training data and Kullback-Leibler (KL) divergence between the predictions from the teacher model and the predictions from the student model; and   adjust weights of the student model, based on the distillation loss for the response-based knowledge distillation technique.   
     
     
         13 . The system of  claim 11 , wherein the knowledge output corresponding to the feature-based knowledge distillation technique comprises internal feature representations obtained from one or more intermediate layers of the teacher model, and wherein the processor-executable instructions further cause the processor to:
 calculate distillation loss for the feature-based knowledge distillation technique, using at least one of: a Euclidean distance or cosine similarity between features of the teacher model and the student model, a mean squared error (MSE) loss, or a correlation alignment loss; and   adjust weights of the student model, based on the distillation loss for feature-based knowledge distillation technique.   
     
     
         14 . The system of  claim 11 , wherein the knowledge output corresponding to the relation-based knowledge distillation technique comprises pair-wise relations and group-wise relations between data points obtained from the teacher model, and wherein the processor-executable instructions further cause the processor to:
 calculate distillation loss for the relation-based knowledge distillation technique, by minimizing discrepancy between class relationships learned by the teacher model and the student model respectively; and   adjust weights of the student model, based on the distillation loss for feature-based knowledge distillation technique.   
     
     
         15 . The system of  claim 10 , wherein training the student model comprises: iteratively inputting each of the one or more batches of training data to the student model based on a predefined epoch. 
     
     
         16 . A non-transitory computer-readable medium storing computer-executable instructions for training a student model using a teacher mode, the computer-executable instructions configured for:
 receiving, from a user, a selection of a target knowledge distillation technique from a plurality of knowledge distillation techniques, a teacher model, a student model, and one or more batches of training data;   loading the teacher model on a memory device;   extracting knowledge output from the teacher model for each of the one or more batches of the training data, based on the target knowledge distillation technique, and sequentially storing extracted knowledge output in a knowledge database, wherein the teacher model is a pre-trained model;   Upon extracting, unloading the teacher model from the memory device and loading the student model on the memory device; and   training the student model based on ground-truth labels associated with each of the one or more batches of training data and the knowledge output corresponding to the target knowledge distillation technique, by fetching the knowledge output corresponding to the target knowledge distillation technique and each of the one or more batches of training data from the knowledge database.   
     
     
         17 . The non-transitory computer-readable medium of  claim 16 , wherein the plurality of knowledge distillation techniques comprises: a response-based knowledge distillation technique, a feature-based knowledge distillation technique, and a relation-based knowledge distillation technique. 
     
     
         18 . The non-transitory computer-readable medium of  claim 17 , wherein the knowledge output corresponding to the response-based knowledge distillation technique comprises soft targets obtained from a final output layer of the teacher model. 
     
     
         19 . The non-transitory computer-readable medium of  claim 18 , wherein the computer-executable instructions are further configured for:
 calculating distillation loss for the response-based knowledge distillation technique, using at least one of: a cross-entropy loss on the ground-truth labels associated with the training data and Kullback-Leibler (KL) divergence between the predictions from the teacher model and the predictions from the student model; and   adjusting weights of the student model, based on the distillation loss for the response-based knowledge distillation technique.   
     
     
         20 . The non-transitory computer-readable medium of  claim 17 , wherein the knowledge output corresponding to the feature-based knowledge distillation technique comprises internal feature representations obtained from one or more intermediate layers of the teacher model.