starting training the model again for updated model env

createPareto.py

@@ -1,6 +1,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+'''
 accuracy = np.array([99.45, 99.45, 98.792, 97.55, 9.45])
 size = np.array([1.0, 1.0, 0.525, 0.45, 0.125])
 
@@ -10,16 +11,18 @@
 plt.xlabel('Size of model (% of largest quantized model)')
 plt.ylabel('Accuracy (MNIST)')
 plt.savefig('size_vs_accuracy.png')
+'''
 
 import pandas as pd
 
-weights = [[0.9, 0.1], [0.75, 0.25], [0.5, 0.5], [0.25, 0.75], [0.1, 0.9]] # do not use 0 because obviously 1-bit for the area part
-colors = ['g', 'b', 'r', 'm', 'k']
+weights = [[1.0, 0.0], [0.8, 0.2]] # do not use 0 because obviously 1-bit for the area part
+#colors = ['g', 'b', 'r', 'm', 'k']
+colors = ['g', 'b']
 plt.figure()
 
 for i, weight in enumerate(weights):
 	agent = pd.read_csv(f'agent_{weights[i][0]}_{weights[i][1]}.monitor.csv', skiprows=0, header = 1)
-	plt.plot(np.arange(100), agent['r'], c = colors[i], label = f'{weights[i][0]}_{weights[i][1]}')
+	plt.plot(np.arange(np.shape(agent)[0]), agent['r'], c = colors[i], label = f'{weights[i][0]}_{weights[i][1]}')
 plt.title('Reward for each agent per episode')
 plt.legend()
 plt.savefig('agents_monitor.png')

train/callbacks/StopTrainingOnNoImprovementCallback.py

@@ -14,7 +14,7 @@ def _on_step(self) -> bool:
         if self.n_calls % self.check_freq == 0:
             # Retrieve training rewards
             training_rewards = self.locals['rewards']
-            mean_reward = np.mean(training_rewards[-100])
+            mean_reward = np.mean(training_rewards[-100:])
 
             if self.verbose > 0:
                 print(f"Mean reward: {mean_reward:.2f} - Best mean reward: {self.best_mean_reward:.2f}")

train/env/ModelEnv.py

@@ -105,7 +105,7 @@ def build_index(self):
         self.model = preprocess_for_quantize(self.model)
         measure_model(self.model, self.model_config['center_crop_shape'], 
                     self.model_config['center_crop_shape'], self.finetuner.in_channels)
-        
+
         self.quantizable_idx = []
         self.layer_types = []
         self.num_quant_acts = 0
@@ -218,9 +218,9 @@ def reset(self, seed = None, option = None):
         return obs, info
 
     def step(self, action):
+
         action = self.get_action(action)
         self.strategy.append(action)
-        self.num_actions += 1
 
         # if not all activations have been quantized
         if self.cur_ind < self.num_quant_acts:
@@ -233,50 +233,44 @@ def step(self, action):
         # if activations have been quantized, quantize outputs and handle residuals
         if self.cur_ind == self.num_quant_acts:
             self.model = self.quantizer.quantize_output(self.model)
-            self.update_index()
-
             self.model = self.quantizer.handle_residuals(self.model)
-            self.update_index()        
-
+
+            # build index again, because quantize output and handle residuals can insert quantizers
+            self.update_index()
+
         if self.cur_ind >= self.num_quant_acts:
             self.model = self.quantizer.quantize_layer(self.model, 
                                                        self.index_to_quantize, 
                                                        int(action[0]))
             # build index again, because quanize layers can insert quantizers
             self.update_index()
 
-        if self.is_final_layer():
-            print(self.strategy)
-            self.quantizer.finalize(self.model)
-
         self.finetuner.model = self.model
         self.finetuner.model.to(self.finetuner.device)
 
-        # accuracy before finetuning
-        acc = self.finetuner.validate(eval = False)
-
-        # can only calibrate if residuals are handled
-        if self.cur_ind >= self.num_quant_acts:
-            self.finetuner.calibrate()
+        self.finetuner.calibrate()
 
-        if self.num_actions % self.args.finetune_every == 0 or self.is_final_layer():
+        # finetune only after all layers have been quantized
+        if self.is_final_layer():
+            self.quantizer.finalize(self.model)
+            self.finetuner.validate()
             self.finetuner.init_finetuning_optim()
             self.finetuner.init_loss()
+            self.finetuner.finetune()
+
+        acc = self.finetuner.validate(eval = False)
 
-            # if it is final layer train for 10 epochs
-            if self.is_final_layer():
-                self.finetuner.finetuning_epochs = 10
-
+        if acc < 25.0:
+            self.finetuner.init_finetuning_optim()
+            self.finetuner.init_loss
             self.finetuner.finetune()
-            self.num_actions = 0
 
         self.action_running_mean = ((action[0]) / (self.max_bit) + (self.cur_ind) * self.action_running_mean) / (self.cur_ind + 1)
         reward = self.reward(acc, action[0])
 
         if self.is_final_layer():
             obs = self.layer_embedding[self.cur_ind, :].copy()
             terminated = True
-            acc = self.finetuner.validate(eval = False)
             info = {"accuracy": acc}
             print(f'Accuracy: {acc}, Size: {self.action_running_mean}')
             return obs, reward, terminated, False, info
@@ -285,8 +279,10 @@ def step(self, action):
         self.cur_ind += 1 
         self.index_to_quantize = self.quantizable_idx[self.cur_ind]
         obs = self.layer_embedding[self.cur_ind, :].copy()
+        if acc < 25.0:
+            acc = self.finetuner.validate(eval = False)
+
         self.prev_acc = acc
-        acc = self.finetuner.validate(eval = False)
         info = {"accuracy": acc}
         return obs, reward, terminated, False, info
 

train/quantizer/Quantizer.py

@@ -504,7 +504,7 @@ def quantize_layer(self,
                         node,
                         rewriters,
                         is_sign_preserving=isinstance(module, SIGN_PRESERVING_MODULES),
-                        quant_identity_map=quant_identity_map,
+                        quant_identity_map=output_quant_identity_map,
                         quant_act_map=quant_act_map,
                         unsigned_act_tuple=unsigned_act_tuple)