def evaluation(model): def my_forward_eval(x, target): out_ = model.forward(x) return out_, model.loss(out_, target) model.eval() accuracy_train, ave_loss_train = compute_acc_loss( my_forward_eval, self.train_loader) print('\tnested train loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss_train, accuracy_train)) # rec.record('train_nested', [ave_loss, accuracy, training_time, step + 1]) accuracy_test, ave_loss_test = compute_acc_loss( my_forward_eval, self.test_loader) model.train() print('\tnested test loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss_test, accuracy_test)) # rec.record('test_nested', [ave_loss, accuracy, training_time, step + 1]) # model.lc_train() return { 'nested_train_loss': ave_loss_train, 'nested_train_acc': accuracy_train, 'nested_test_loss': ave_loss_test, 'nested_acc': accuracy_test }
def finetuning(config): all_start_time = time.time() my_params = filter(lambda p: p.requires_grad, self.model.parameters()) optimizer = torch.optim.SGD(my_params, config['lr'], momentum=config['momentum'], nesterov=True) epoch_time = AverageMeter() lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=config['lr_decay']) train_info = {} test_info = {} from threading import Thread import asyncio def start_loop(loop): asyncio.set_event_loop(loop) loop.run_forever() new_loop = asyncio.new_event_loop() t = Thread(target=start_loop, args=(new_loop, )) t.start() for epoch in range(0, config['epochs']): start_time = time.time() avg_loss_ = AverageMeter() for x, target in self.train_loader: optimizer.zero_grad() x, target = x.cuda(), target.cuda() loss = self.model.loss(self.model(x), target) loss.backward() avg_loss_.update(loss.item()) optimizer.step() end_time = time.time() training_time = end_time - all_start_time epoch_time.update(end_time - start_time) print( "Epoch {0} finished in {1.val:.3f}s (avg: {1.avg:.3f}s). Training for {2}" .format(epoch, epoch_time, format_time(end_time - all_start_time))) print('AVG train loss {0.avg:.6f}'.format(avg_loss_)) print("\tLR: {:.4e}".format(lr_scheduler.get_lr()[0])) lr_scheduler.step() if (epoch + 1) % config['print_freq'] == 0: self.model.eval() accuracy, ave_loss = compute_acc_loss( my_forward_eval, self.train_loader) train_info[epoch + 1] = [ave_loss, accuracy, training_time] print('\ttrain loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) accuracy, ave_loss = compute_acc_loss( my_forward_eval, self.test_loader) test_info[epoch + 1] = [ave_loss, accuracy, training_time] print('\ttest loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) self.model.train() to_save = {} to_save['config'] = config to_save['optimizer_state'] = optimizer.state_dict() to_save['model_state'] = self.model.state_dict() to_save['training_time'] = training_time to_save['traing_info'] = train_info to_save['test_info'] = test_info to_save['current_epoch'] = epoch to_save['compression_stats'] = compression_stats async def actual_save(): # TODO: make better saves, 1) mv file as backup, 2) save new data 3) delte bk torch.save( to_save, f'results/{self.name}_ft_{config["tag"]}.th') asyncio.run_coroutine_threadsafe(actual_save(), new_loop) async def last_task(): print("Async file saving has been finished.") new_loop.stop() asyncio.run_coroutine_threadsafe(last_task(), new_loop)
def finetune_setup(self, tag_of_lc_model, c_step_config): exp_run_details = torch.load( f"results/{self.name}_lc_{tag_of_lc_model}.th", map_location="cpu") # despite the 's' at the end, there is only one model state, the last model_state_to_load = exp_run_details['model_states'] last_lc_it = exp_run_details['last_step_number'] model = add_flops_counting_methods(self.model) model.start_flops_count() for x, target in self.train_loader: _ = model(x[None, 0].cuda()) break all_flops = model.compute_average_flops_cost() model.stop_flops_count() self.model = self.model.cpu() all_params = count_params(self.model) compression_info = {} for task_name, infos in exp_run_details[ 'compression_tasks_info'].items(): compression_info[task_name] = infos[last_lc_it] # compression_infos = exp_run_details['compression_tasks_info'][last_lc_it] print(model_state_to_load.keys()) for key in list(model_state_to_load.keys()): if key.startswith("lc_param_list"): del model_state_to_load[key] del exp_run_details import gc gc.collect() print(model_state_to_load.keys()) self.model.load_state_dict(model_state_to_load) print("model has been sucessfully loaded") for i, module in enumerate([ x for x in self.model.modules() if isinstance(x, nn.Conv2d) or isinstance(x, nn.Linear) ]): module.selected_rank_ = compression_info[f"task_{i}"][ 'selected_rank'] print(module.selected_rank_) reparametrize_low_rank(self.model) print(self.model) self.model = self.model.to(self.device) print( "Low rank layers of the model has been successfully reparameterized with sequence of full-rank matrices." ) def my_forward_eval(x, target): out_ = self.model.forward(x) return out_, self.model.loss(out_, target) self.model.eval() accuracy_train, ave_loss_train = compute_acc_loss( my_forward_eval, self.train_loader) print('\tBefore finetuning, the train loss: {:.6f}, accuracy: {:.4f}'. format(ave_loss_train, accuracy_train)) # rec.record('train_nested', [ave_loss, accuracy, training_time, step + 1]) accuracy_test, ave_loss_test = compute_acc_loss( my_forward_eval, self.test_loader) self.model.train() print('\tBefore finetuning, the test loss: {:.6f}, accuracy: {:.4f}'. format(ave_loss_test, accuracy_test)) model = add_flops_counting_methods(self.model) model.start_flops_count() for x, target in self.train_loader: _ = model(x[None, 0].cuda()) break compressed_flops = model.compute_average_flops_cost() model.stop_flops_count() self.model = self.model.cpu() compressed_params = count_params(model) self.model = self.model.to(self.device) print('The number of FLOPS in original model', all_flops) print('The number of params in original model:', all_params) print('The number of FLOPS in this model', compressed_flops) print('The number of params in this model:', compressed_params) flops_rho = all_flops[0] / compressed_flops[0] storage_rho = all_params / compressed_params print(f'FLOPS ρ={flops_rho:.3f}; STORAGE ρ={storage_rho:.3f};') compression_stats = { 'original_flops': all_flops, 'compressed_flops': compressed_flops, 'flops_rho': flops_rho, 'storage_rho': storage_rho } def finetuning(config): all_start_time = time.time() my_params = filter(lambda p: p.requires_grad, self.model.parameters()) optimizer = torch.optim.SGD(my_params, config['lr'], momentum=config['momentum'], nesterov=True) epoch_time = AverageMeter() lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=config['lr_decay']) train_info = {} test_info = {} from threading import Thread import asyncio def start_loop(loop): asyncio.set_event_loop(loop) loop.run_forever() new_loop = asyncio.new_event_loop() t = Thread(target=start_loop, args=(new_loop, )) t.start() for epoch in range(0, config['epochs']): start_time = time.time() avg_loss_ = AverageMeter() for x, target in self.train_loader: optimizer.zero_grad() x, target = x.cuda(), target.cuda() loss = self.model.loss(self.model(x), target) loss.backward() avg_loss_.update(loss.item()) optimizer.step() end_time = time.time() training_time = end_time - all_start_time epoch_time.update(end_time - start_time) print( "Epoch {0} finished in {1.val:.3f}s (avg: {1.avg:.3f}s). Training for {2}" .format(epoch, epoch_time, format_time(end_time - all_start_time))) print('AVG train loss {0.avg:.6f}'.format(avg_loss_)) print("\tLR: {:.4e}".format(lr_scheduler.get_lr()[0])) lr_scheduler.step() if (epoch + 1) % config['print_freq'] == 0: self.model.eval() accuracy, ave_loss = compute_acc_loss( my_forward_eval, self.train_loader) train_info[epoch + 1] = [ave_loss, accuracy, training_time] print('\ttrain loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) accuracy, ave_loss = compute_acc_loss( my_forward_eval, self.test_loader) test_info[epoch + 1] = [ave_loss, accuracy, training_time] print('\ttest loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) self.model.train() to_save = {} to_save['config'] = config to_save['optimizer_state'] = optimizer.state_dict() to_save['model_state'] = self.model.state_dict() to_save['training_time'] = training_time to_save['traing_info'] = train_info to_save['test_info'] = test_info to_save['current_epoch'] = epoch to_save['compression_stats'] = compression_stats async def actual_save(): # TODO: make better saves, 1) mv file as backup, 2) save new data 3) delte bk torch.save( to_save, f'results/{self.name}_ft_{config["tag"]}.th') asyncio.run_coroutine_threadsafe(actual_save(), new_loop) async def last_task(): print("Async file saving has been finished.") new_loop.stop() asyncio.run_coroutine_threadsafe(last_task(), new_loop) return finetuning
def l_step_optimization(model, lc_penalty, step, config): all_start_time = config['all_start_time'] lr_scheduler = None my_params = filter(lambda p: p.requires_grad, model.parameters()) learning_rate = config['lr'] if config['lr_decay_mode'] == 'after_l': learning_rate *= (config['lr_decay']**step) print(f"Current LR={learning_rate}") def constract_my_forward_lc_eval(lc_penalty): pen = lc_penalty() def my_forward_lc_eval(x, target): out_ = model.forward(x) return out_, model.loss(out_, target) + pen return my_forward_lc_eval optimizer = torch.optim.SGD(my_params, learning_rate, momentum=config['momentum'], nesterov=True) if config['lr_decay_mode'] == 'restart_on_l': lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=config['lr_decay']) if 'lr_trick' in config: l_trick_value = 0.1 print( 'LR trick in play. first epoch is trained with LR of {:.4e}' .format(config['lr'] * l_trick_value)) for param_group in optimizer.param_groups: param_group['lr'] = config['lr'] * l_trick_value # TODO: revert back the lr_trick? epochs_in_this_it = config['epochs'] if step > 0 else \ config['first_mu_epochs'] if 'first_mu_epochs' in config else config['epochs'] print('Epochs in this iteration is :', epochs_in_this_it) print('Epochs in this iteration is :', epochs_in_this_it) model.eval() lc_evaluator = constract_my_forward_lc_eval(lc_penalty) accuracy, ave_loss = compute_acc_loss(lc_evaluator, self.train_loader) print('\ttrain loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) accuracy, ave_loss = compute_acc_loss(lc_evaluator, self.test_loader) print('\ttest loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) model.train() epoch_time = AverageMeter() rec = Recorder() # avg_epoch_losses = [] for epoch in range(epochs_in_this_it): start_time = time.time() avg_loss_ = AverageMeter() for x, target in self.train_loader: optimizer.zero_grad() x, target = x.cuda(), target.cuda(non_blocking=True) loss = model.loss(model(x), target) + lc_penalty() avg_loss_.update(loss.item()) loss.backward() optimizer.step() end_time = time.time() training_time = end_time - all_start_time epoch_time.update(end_time - start_time) print( "LC step {0}, Epoch {1} finished in {2.val:.3f}s (avg: {2.avg:.3f}s). Training for {3}" .format(step, epoch, epoch_time, format_time(end_time - all_start_time))) print('AVG train loss {0.avg:.6f}'.format(avg_loss_)) rec.record('average_loss_per_epoch', avg_loss_) if (epoch + 1) % config['print_freq'] == 0: model.eval() lc_evaluator = constract_my_forward_lc_eval(lc_penalty) accuracy, ave_loss = compute_acc_loss( lc_evaluator, self.train_loader) rec.record('train', [ ave_loss, accuracy, training_time, step + 1, epoch + 1 ]) print('\ttrain loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) accuracy, ave_loss = compute_acc_loss( lc_evaluator, self.test_loader) rec.record('test', [ ave_loss, accuracy, training_time, step + 1, epoch + 1 ]) print('\ttest loss: {:.6f}, accuracy: {:.4f}'.format( ave_loss, accuracy)) model.train() if config['lr_decay_mode'] == 'restart_on_l': print("\told LR: {:.4e}".format( optimizer.param_groups[0]['lr'])) lr_scheduler.step() print("\tnew LR: {:.4e}".format( optimizer.param_groups[0]['lr'])) else: print("\tLR: {:.4e}".format(learning_rate)) info = { 'train': rec.train, 'test': rec.test, 'average_loss_per_train_epoch': rec.average_loss_per_epoch } return info