From d5bb7033f7d270cf75ad0c5a573f0dd7f478af19 Mon Sep 17 00:00:00 2001 From: Het-Shah Date: Fri, 11 Jun 2021 16:24:48 +0530 Subject: [PATCH 1/2] Correct tests --- .travis.yml | 1 - KD_Lib/KD/common/base_class.py | 34 ++++++++++++++++----------------- KD_Lib/KD/vision/DML/dml.py | 4 +++- tests/test_KD_Lib.py | 35 +++++++++++++++++++++++++++------- 4 files changed, 47 insertions(+), 27 deletions(-) diff --git a/.travis.yml b/.travis.yml index 580eec0d..b7894402 100644 --- a/.travis.yml +++ b/.travis.yml @@ -13,7 +13,6 @@ install: - pip install -U tox-travis codecov black - python setup.py install - jobs: include: # Deploy Documentation diff --git a/KD_Lib/KD/common/base_class.py b/KD_Lib/KD/common/base_class.py index 4c0328e2..9d74220d 100755 --- a/KD_Lib/KD/common/base_class.py +++ b/KD_Lib/KD/common/base_class.py @@ -53,28 +53,26 @@ def __init__( if self.log: self.writer = SummaryWriter(logdir) - try: - torch.Tensor(0).to(device) - self.device = device - except: - print( - "Either an invalid device or CUDA is not available. Defaulting to CPU." - ) + if device == "cpu": self.device = torch.device("cpu") + elif device == "cuda": + if torch.cuda.is_available(): + self.device = torch.device("cuda") + else: + print( + "Either an invalid device or CUDA is not available. Defaulting to CPU." + ) + self.device = torch.device("cpu") - try: + if teacher_model: self.teacher_model = teacher_model.to(self.device) - except: + else: print("Warning!!! Teacher is NONE.") + self.student_model = student_model.to(self.device) - try: - self.loss_fn = loss_fn.to(self.device) - self.ce_fn = nn.CrossEntropyLoss().to(self.device) - except: - self.loss_fn = loss_fn - self.ce_fn = nn.CrossEntropyLoss() - print("Warning: Loss Function can't be moved to device.") - + self.loss_fn = loss_fn.to(self.device) + self.ce_fn = nn.CrossEntropyLoss().to(self.device) + def train_teacher( self, epochs=20, @@ -208,7 +206,7 @@ def _train_student( epoch_acc = correct / length_of_dataset - epoch_val_acc = self.evaluate(teacher=False) + _, epoch_val_acc = self._evaluate_model(self.student_model, verbose=True) if epoch_val_acc > best_acc: best_acc = epoch_val_acc diff --git a/KD_Lib/KD/vision/DML/dml.py b/KD_Lib/KD/vision/DML/dml.py index cd47ad9c..20f7dd84 100755 --- a/KD_Lib/KD/vision/DML/dml.py +++ b/KD_Lib/KD/vision/DML/dml.py @@ -171,7 +171,9 @@ def _evaluate_model(self, model, verbose=True): if verbose: print(f"Accuracy: {correct/length_of_dataset}") - return outputs + + epoch_val_acc = correct/length_of_dataset + return outputs, epoch_val_acc def evaluate(self): """ diff --git a/tests/test_KD_Lib.py b/tests/test_KD_Lib.py index 4989fb2a..f3f4a4cd 100644 --- a/tests/test_KD_Lib.py +++ b/tests/test_KD_Lib.py @@ -23,6 +23,7 @@ ProbShift, LabelSmoothReg, DML, + BaseClass ) from KD_Lib.models import ( @@ -94,6 +95,10 @@ def test_resnet(): ResNet50(params) ResNet101(params) ResNet152(params) + ResNet34(params, att=True) + ResNet34(params, mean=True) + ResNet101(params, att=True) + ResNet101(params, mean=True) def test_attention_model(): @@ -159,6 +164,22 @@ def test_LSTMNet(): # Strategy TESTS # +def test_BaseClass() + teac = Shallow(hidden_size=400) + stud = Shallow(hidden_size=100) + + t_optimizer = optim.SGD(teac.parameters(), 0.01) + s_optimizer = optim.SGD(stud.parameters(), 0.01) + + distiller = BaseClass( + teac, stud, train_loader, test_loader, t_optimizer, s_optimizer, log=True + ) + + distiller.train_teacher(epochs=1, plot_losses=True, save_model=True) + distiller.train_student(epochs=1, plot_losses=True, save_model=True) + distiller.evaluate(teacher=False) + distiller.get_parameters() + def test_VanillaKD(): teac = Shallow(hidden_size=400) @@ -168,11 +189,11 @@ def test_VanillaKD(): s_optimizer = optim.SGD(stud.parameters(), 0.01) distiller = VanillaKD( - teac, stud, train_loader, test_loader, t_optimizer, s_optimizer + teac, stud, train_loader, test_loader, t_optimizer, s_optimizer, log=True ) - distiller.train_teacher(epochs=1, plot_losses=False, save_model=False) - distiller.train_student(epochs=1, plot_losses=False, save_model=False) + distiller.train_teacher(epochs=1, plot_losses=True, save_model=True) + distiller.train_student(epochs=1, plot_losses=True, save_model=True) distiller.evaluate(teacher=False) distiller.get_parameters() @@ -289,8 +310,8 @@ def test_SelfTraining(): def test_mean_teacher(): - teacher_params = [4, 4, 8, 4, 4] - student_params = [4, 4, 4, 4, 4] + teacher_params = [16, 16, 32, 16, 16] + student_params = [16, 16, 16, 16, 16] teacher_model = ResNet50(teacher_params, 1, 10, mean=True) student_model = ResNet18(student_params, 1, 10, mean=True) @@ -488,7 +509,7 @@ def test_lottery_tickets(): teacher_params = [4, 4, 8, 4, 4] teacher_model = ResNet50(teacher_params, 1, 10, True) pruner = Lottery_Tickets_Pruner(teacher_model, train_loader, test_loader) - pruner.prune(num_iterations=1, train_iterations=1, valid_freq=1, print_freq=1) + pruner.prune(num_iterations=2, train_iterations=2, valid_freq=1, print_freq=1) # @@ -539,6 +560,6 @@ def test_qat_quantization(): model.fc.out_features = 10 optimizer = torch.optim.Adam(model.parameters()) quantizer = QAT_Quantizer(model, cifar_trainloader, cifar_testloader, optimizer) - quantized_model = quantizer.quantize(1, 1, -1, -1) + quantized_model = quantizer.quantize(1, 1, 1, 1) quantizer.get_model_sizes() quantizer.get_performance_statistics() From 1a1b0ecb85b85c239c3f8860be5743574636567d Mon Sep 17 00:00:00 2001 From: Het-Shah Date: Fri, 11 Jun 2021 16:47:59 +0530 Subject: [PATCH 2/2] Update tests and apply black --- .github/workflows/python-package-test.yml | 5 ----- KD_Lib/KD/common/base_class.py | 14 +++++++------- KD_Lib/KD/vision/DML/dml.py | 2 +- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/.github/workflows/python-package-test.yml b/.github/workflows/python-package-test.yml index 9fae0f4b..0ea735dd 100644 --- a/.github/workflows/python-package-test.yml +++ b/.github/workflows/python-package-test.yml @@ -33,11 +33,6 @@ jobs: pip install build - name: Build package run: python -m build - - name: Black - run: | - # stop the build if there are Python syntax errors or undefined names - black --check KD_Lib - black --check tests - name: Test with pytest run: | pytest diff --git a/KD_Lib/KD/common/base_class.py b/KD_Lib/KD/common/base_class.py index 9d74220d..74ed175b 100755 --- a/KD_Lib/KD/common/base_class.py +++ b/KD_Lib/KD/common/base_class.py @@ -68,11 +68,11 @@ def __init__( self.teacher_model = teacher_model.to(self.device) else: print("Warning!!! Teacher is NONE.") - + self.student_model = student_model.to(self.device) self.loss_fn = loss_fn.to(self.device) self.ce_fn = nn.CrossEntropyLoss().to(self.device) - + def train_teacher( self, epochs=20, @@ -140,7 +140,7 @@ def train_teacher( ) loss_arr.append(epoch_loss) - print(f"Epoch: {ep+1}, Loss: {epoch_loss}, Accuracy: {epoch_acc}") + print("Epoch: {}, Loss: {}, Accuracy: {}".format(ep+1, epoch_loss, epoch_acc)) self.post_epoch_call(ep) @@ -222,7 +222,7 @@ def _train_student( ) loss_arr.append(epoch_loss) - print(f"Epoch: {ep+1}, Loss: {epoch_loss}, Accuracy: {epoch_acc}") + print("Epoch: {}, Loss: {}, Accuracy: {}".format(ep+1, epoch_loss, epoch_acc)) self.student_model.load_state_dict(self.best_student_model_weights) if save_model: @@ -288,7 +288,7 @@ def _evaluate_model(self, model, verbose=True): if verbose: print("-" * 80) - print(f"Validation Accuracy: {accuracy}") + print("Validation Accuracy: {}".format(accuracy)) return outputs, accuracy def evaluate(self, teacher=False): @@ -313,8 +313,8 @@ def get_parameters(self): student_params = sum(p.numel() for p in self.student_model.parameters()) print("-" * 80) - print(f"Total parameters for the teacher network are: {teacher_params}") - print(f"Total parameters for the student network are: {student_params}") + print("Total parameters for the teacher network are: {}".format(teacher_params)) + print("Total parameters for the student network are: {}".format(student_params)) def post_epoch_call(self, epoch): """ diff --git a/KD_Lib/KD/vision/DML/dml.py b/KD_Lib/KD/vision/DML/dml.py index 20f7dd84..8b37347d 100755 --- a/KD_Lib/KD/vision/DML/dml.py +++ b/KD_Lib/KD/vision/DML/dml.py @@ -172,7 +172,7 @@ def _evaluate_model(self, model, verbose=True): if verbose: print(f"Accuracy: {correct/length_of_dataset}") - epoch_val_acc = correct/length_of_dataset + epoch_val_acc = correct / length_of_dataset return outputs, epoch_val_acc def evaluate(self):