def train_auto_compressor(self): using_fm_reconstruction = self.args.amc_prune_method == 'fm-reconstruction' fixed_subset, sequential = (using_fm_reconstruction, using_fm_reconstruction) msglogger.info("AMC: fixed_subset=%s\tsequential=%s" % (fixed_subset, sequential)) train_loader, val_loader, test_loader = classifier.load_data( self.args, fixed_subset, sequential) self.args.display_confusion = False validate_fn = partial(classifier.test, test_loader=val_loader, criterion=self.criterion, loggers=self.pylogger, args=self.args, activations_collectors=None) train_fn = partial(classifier.train, train_loader=train_loader, criterion=self.criterion, loggers=self.pylogger, args=self.args) save_checkpoint_fn = partial(apputils.save_checkpoint, arch=self.args.arch, dir=msglogger.logdir) optimizer_data = { 'lr': self.args.lr, 'momentum': self.args.momentum, 'weight_decay': self.args.weight_decay } return train_auto_compressor(self.model, self.args, optimizer_data, validate_fn, save_checkpoint_fn, train_fn)
def greedy(model, criterion, optimizer, loggers, args): train_loader, val_loader, test_loader = classifier.load_data(args) test_fn = partial(classifier.test, test_loader=test_loader, criterion=criterion, loggers=loggers, args=args, activations_collectors=None) train_fn = partial(classifier.train, train_loader=train_loader, criterion=criterion, args=args) assert args.greedy_target_density is not None distiller.pruning.greedy_filter_pruning.greedy_pruner(model, args, args.greedy_target_density, args.greedy_pruning_step, test_fn, train_fn)
def load_test_data(args): test_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True) return test_loader
q_weight_name = module_name + ".weight" if q_weight_name == name_in: return module if __name__ == "__main__": print("START") args = get_default_args() # args.epochs = float('inf') # hack for args parsing so there's no error in epochs cc = classifier.ClassifierCompressor(args, script_dir=os.path.dirname(__file__)) args = deepcopy( cc.args ) # Get back args after modifications in ClassifierCompressor.__init__ eval_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True) # logging logging.getLogger().setLevel(logging.WARNING) msglogger = logging.getLogger(__name__) msglogger.setLevel(logging.INFO) def test_fn(model): top1, top5, losses = classifier.test(eval_data_loader, model, cc.criterion, [cc.tflogger, cc.pylogger], None, args) # pdb.set_trace() return top1, top5, losses
def finetune_directory(ft_dir, stats_file, app_args, cleanup_ft_dir=False, checkpoints=None): """Fine tune all the checkpoint files we find in the immediate-directory specified. For each checkpoint file we find, we create and queue a FinetuningTask. A FinetuningProcess will pickup the FinetuningTask and process it. """ print("Fine-tuning directory %s" % ft_dir) if not checkpoints: # Get a list of the checkpoint files checkpoints = glob.glob(os.path.join(ft_dir, "*checkpoint.pth.tar")) assert checkpoints # We create a subdirectory, where we will write all of our output ft_output_dir = os.path.join(ft_dir, "ft") os.makedirs(ft_output_dir, exist_ok=True) print("Writing results to directory %s" % ft_output_dir) app_args.output_dir = ft_output_dir # Multi-process queues tasks = multiprocessing.JoinableQueue() results = multiprocessing.Queue() # Create and launch the fine-tuning processes processes = [] n_processes = min(app_args.processes, len(checkpoints)) for i in range(n_processes): # Pre-load the data-loaders of each fine-tuning process once app = classifier.ClassifierCompressor( app_args, script_dir=os.path.dirname(__file__)) data_loader = classifier.load_data(app.args) # Delete log directories shutil.rmtree(app.logdir) processes.append(FinetuningProcess(tasks, results, data_loader)) # Start the process processes[-1].start() n_gpus = torch.cuda.device_count() # Enqueue all of the fine-tuning tasks for (instance, ckpt_file) in enumerate(checkpoints): tasks.put(FinetuningTask(args=(ckpt_file, instance % n_gpus, app_args))) # Push an end-of-tasks marker for i in range(len(processes)): tasks.put(None) # Wait until all tasks finish tasks.join() # Start printing results results_dict = OrderedDict() while not results.empty(): result = results.get() results_dict[result[0]] = result[1] # Read the results of the AMC experiment (we'll want to use some of the data) import pandas as pd df = pd.read_csv(os.path.join(ft_dir, "amc.csv")) assert len(results_dict) > 0 # Log some info for each checkpoint for ckpt_name in sorted(results_dict.keys()): net_search_results = df[df["ckpt_name"] == ckpt_name[:-len("_checkpoint.pth.tar")]] search_top1 = net_search_results["top1"].iloc[0] normalized_macs = net_search_results["normalized_macs"].iloc[0] log_entry = (ft_output_dir, ckpt_name, normalized_macs, search_top1, *results_dict[ckpt_name]) print("%s <> %s: %.2f %.2f %.2f %.2f %.2f" % log_entry) stats_file.add_record(log_entry) if cleanup_ft_dir: # cleanup: remove the "ft" directory shutil.rmtree(ft_output_dir)
msglogger.info('best_overrides_dict: %s' % best_overrides_dict) msglogger.info('Best score: %f' % eval_fn(quantizer.model)) return model, best_overrides_dict if __name__ == "__main__": args = get_default_args() args.epochs = float( 'inf') # hack for args parsing so there's no error in epochs cc = classifier.ClassifierCompressor(args, script_dir=os.path.dirname(__file__)) args = deepcopy( cc.args ) # Get back args after modifications in ClassifierCompressor.__init__ eval_data_loader = classifier.load_data(args, load_train=False, load_val=False) # quant calibration dataloader: args.effective_test_size = args.qe_calib_portion args.batch_size = args.qe_calib_batchsize calib_data_loader = classifier.load_data(args, fixed_subset=True, load_train=False, load_val=False) # logging logging.getLogger().setLevel(logging.WARNING) msglogger = logging.getLogger(__name__) msglogger.setLevel(logging.INFO) def test_fn(model):
def image_classifier_ptq_lapq(model, criterion, loggers, args): # data loader function for splitting the validation set. args = deepcopy(args) effective_test_size_bak = args.effective_test_size args.effective_test_size = args.lapq_eval_size eval_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True, fixed_subset=True) args.effective_test_size = effective_test_size_bak test_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True) model = model.eval() device = next(model.parameters()).device if args.lapq_eval_memoize_dataloader: images_batches = [] targets_batches = [] for images, targets in eval_data_loader: images_batches.append(images.to(device)) targets_batches.append(targets.to(device)) memoized_data_loader = [(torch.cat(images_batches), torch.cat(targets_batches))] else: memoized_data_loader = None def eval_fn(model): if memoized_data_loader: loss = 0 for images, targets in memoized_data_loader: outputs = model(images) loss += criterion(outputs, targets).item() loss = loss / len(memoized_data_loader) else: _, _, loss = classifier.test(eval_data_loader, model, criterion, loggers, None, args) return loss def test_fn(model): top1, top5, loss = classifier.test(test_data_loader, model, criterion, loggers, None, args) return OrderedDict([('top-1', top1), ('top-5', top5), ('loss', loss)]) args.device = device if args.resumed_checkpoint_path: args.load_model_path = args.resumed_checkpoint_path if args.load_model_path: msglogger.info("Loading checkpoint from %s" % args.load_model_path) model = apputils.load_lean_checkpoint(model, args.load_model_path, model_device=args.device) quantizer = distiller.quantization.PostTrainLinearQuantizer.from_args(model, args) dummy_input = torch.rand(*model.input_shape, device=args.device) model, qp_dict = lapq.ptq_coordinate_search(quantizer, dummy_input, eval_fn, test_fn=test_fn, **lapq.cmdline_args_to_dict(args)) results = test_fn(quantizer.model) msglogger.info("Arch: %s \tTest: \t top1 = %.3f \t top5 = %.3f \t loss = %.3f" % (args.arch, results['top-1'], results['top-5'], results['loss'])) distiller.yaml_ordered_save('%s.quant_params_dict.yaml' % args.arch, qp_dict) distiller.apputils.save_checkpoint(0, args.arch, model, extras={'top1': results['top-1'], 'qp_dict': qp_dict}, name=args.name, dir=msglogger.logdir)
if __name__ == "__main__": args = get_default_args() args.epochs = float( 'inf') # hack for args parsing so there's no error in epochs cc = classifier.ClassifierCompressor(args, script_dir=os.path.dirname(__file__)) args = deepcopy(cc.args) effective_test_size_bak = args.effective_test_size args.effective_test_size = args.opt_val_size eval_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True, fixed_subset=True) args.effective_test_size = effective_test_size_bak test_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True) # logging logging.getLogger().setLevel(logging.WARNING) msglogger = logging.getLogger(__name__) msglogger.setLevel(logging.INFO) model = cc.model.eval()