def _remove_bottleneck(self):
replace_layer(self.model, self.sequential, self.target_layer)
self.bn_layer.active = False
def _run_training(self, input_t, target_t):
# Attach layer and train the bottleneck
replace_layer(self.model, self.original_layer, self.sequential)
self._train_bottleneck(input_t, target_t)
replace_layer(self.model, self.sequential, self.original_layer)
def _inject_bottleneck(self):
replace_layer(self.model, self.target_layer, self.sequential)
self.bn_layer.active = True