class Runner(): ''' Handler for complete pre-training progress of upstream models ''' def __init__(self, args, config, dataloader, ckpdir): self.device = torch.device('cuda') if ( args.gpu and torch.cuda.is_available()) else torch.device('cpu') if torch.cuda.is_available(): print('[Runner] - CUDA is available!') self.model_kept = [] self.global_step = 1 self.log = SummaryWriter(ckpdir) self.args = args self.config = config self.dataloader = dataloader self.ckpdir = ckpdir # optimizer self.learning_rate = float(config['optimizer']['learning_rate']) self.warmup_proportion = config['optimizer']['warmup_proportion'] self.gradient_accumulation_steps = config['optimizer'][ 'gradient_accumulation_steps'] self.gradient_clipping = config['optimizer']['gradient_clipping'] # Training details self.apex = config['runner']['apex'] self.total_steps = config['runner']['total_steps'] self.log_step = config['runner']['log_step'] self.save_step = config['runner']['save_step'] self.duo_feature = config['runner']['duo_feature'] self.max_keep = config['runner']['max_keep'] # model self.transformer_config = config['transformer'] self.input_dim = self.transformer_config['input_dim'] self.output_dim = 1025 if self.duo_feature else None # output dim is the same as input dim if not using duo features def set_model(self): print('[Runner] - Initializing Transformer model...') # build the Transformer model with speech prediction head model_config = TransformerConfig(self.config) self.dr = model_config.downsample_rate self.hidden_size = model_config.hidden_size self.model = TransformerForMaskedAcousticModel( model_config, self.input_dim, self.output_dim).to(self.device) self.model.train() if self.args.multi_gpu: self.model = torch.nn.DataParallel(self.model) print('[Runner] - Multi-GPU training Enabled: ' + str(torch.cuda.device_count())) print('[Runner] - Number of parameters: ' + str( sum(p.numel() for p in self.model.parameters() if p.requires_grad))) # Setup optimizer param_optimizer = list(self.model.named_parameters()) no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [{ 'params': [ p for n, p in param_optimizer if not any(nd in n for nd in no_decay) ], 'weight_decay': 0.01 }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] if self.apex: try: from apex.optimizers import FP16_Optimizer from apex.optimizers import FusedAdam except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training." ) optimizer = FusedAdam(optimizer_grouped_parameters, lr=self.learning_rate, bias_correction=False, max_grad_norm=1.0) if self.config['optimizer']['loss_scale'] == 0: self.optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True) else: self.optimizer = FP16_Optimizer( optimizer, static_loss_scale=self.config['optimizer']['loss_scale']) self.warmup_linear = WarmupLinearSchedule( warmup=self.warmup_proportion, t_total=self.total_steps) else: self.optimizer = BertAdam(optimizer_grouped_parameters, lr=self.learning_rate, warmup=self.warmup_proportion, t_total=self.total_steps) def save_model(self, name='states', to_path=None): all_states = { 'SpecHead': self.model.SpecHead.state_dict() if not self.args.multi_gpu else self.model.module.SpecHead.state_dict(), 'Transformer': self.model.Transformer.state_dict() if not self.args.multi_gpu else self.model.module.Transformer.state_dict(), 'Optimizer': self.optimizer.state_dict(), 'Global_step': self.global_step, 'Settings': { 'Config': self.config, 'Paras': self.args, }, } if to_path is None: new_model_path = '{}/{}-{}.ckpt'.format(self.ckpdir, name, self.global_step) else: new_model_path = to_path torch.save(all_states, new_model_path) self.model_kept.append(new_model_path) if len(self.model_kept) >= self.max_keep: os.remove(self.model_kept[0]) self.model_kept.pop(0) def up_sample_frames(self, spec, return_first=False): if len(spec.shape) != 3: spec = spec.unsqueeze(0) assert (len(spec.shape) == 3 ), 'Input should have acoustic feature of shape BxTxD' # spec shape: [batch_size, sequence_length // downsample_rate, output_dim * downsample_rate] spec_flatten = spec.view(spec.shape[0], spec.shape[1] * self.dr, spec.shape[2] // self.dr) if return_first: return spec_flatten[0] return spec_flatten # spec_flatten shape: [batch_size, sequence_length * downsample_rate, output_dim // downsample_rate] def down_sample_frames(self, spec): left_over = spec.shape[1] % self.dr if left_over != 0: spec = spec[:, :-left_over, :] spec_stacked = spec.view(spec.shape[0], spec.shape[1] // self.dr, spec.shape[2] * self.dr) return spec_stacked def process_data(self, spec): """Process training data for the masked acoustic model""" with torch.no_grad(): assert ( len(spec) == 5 ), 'dataloader should return (spec_masked, pos_enc, mask_label, attn_mask, spec_stacked)' # Unpack and Hack bucket: Bucketing should cause acoustic feature to have shape 1xBxTxD' spec_masked = spec[0].squeeze(0) pos_enc = spec[1].squeeze(0) mask_label = spec[2].squeeze(0) attn_mask = spec[3].squeeze(0) spec_stacked = spec[4].squeeze(0) spec_masked = spec_masked.to(device=self.device) if pos_enc.dim() == 3: # pos_enc: (batch_size, seq_len, hidden_size) # GPU memory need (batch_size * seq_len * hidden_size) pos_enc = torch.FloatTensor(pos_enc).to(device=self.device) elif pos_enc.dim() == 2: # pos_enc: (seq_len, hidden_size) # GPU memory only need (seq_len * hidden_size) even after expanded pos_enc = torch.FloatTensor(pos_enc).to( device=self.device).expand(spec_masked.size(0), *pos_enc.size()) mask_label = torch.ByteTensor(mask_label).to(device=self.device) attn_mask = torch.FloatTensor(attn_mask).to(device=self.device) spec_stacked = spec_stacked.to(device=self.device) return spec_masked, pos_enc, mask_label, attn_mask, spec_stacked # (x, pos_enc, mask_label, attention_mask. y) def train(self): ''' Self-Supervised Pre-Training of Transformer Model''' pbar = tqdm(total=self.total_steps) while self.global_step <= self.total_steps: progress = tqdm(self.dataloader, desc="Iteration") step = 0 loss_val = 0 for batch_is_valid, *batch in progress: try: if self.global_step > self.total_steps: break if not batch_is_valid: continue step += 1 spec_masked, pos_enc, mask_label, attn_mask, spec_stacked = self.process_data( batch) loss, pred_spec = self.model(spec_masked, pos_enc, mask_label, attn_mask, spec_stacked) # Accumulate Loss if self.gradient_accumulation_steps > 1: loss = loss / self.gradient_accumulation_steps if self.apex and self.args.multi_gpu: raise NotImplementedError elif self.apex: self.optimizer.backward(loss) elif self.args.multi_gpu: loss = loss.sum() loss.backward() else: loss.backward() loss_val += loss.item() # Update if (step + 1) % self.gradient_accumulation_steps == 0: if self.apex: # modify learning rate with special warm up BERT uses # if conifg.apex is False, BertAdam is used and handles this automatically lr_this_step = self.learning_rate * self.warmup_linear.get_lr( self.global_step, self.warmup_proportion) for param_group in self.optimizer.param_groups: param_group['lr'] = lr_this_step # Step grad_norm = torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.gradient_clipping) if math.isnan(grad_norm): print( '[Runner] - Error : grad norm is NaN @ step ' + str(self.global_step)) else: self.optimizer.step() self.optimizer.zero_grad() if self.global_step % self.log_step == 0: # Log self.log.add_scalar('lr', self.optimizer.get_lr()[0], self.global_step) self.log.add_scalar('loss', (loss_val), self.global_step) self.log.add_scalar('gradient norm', grad_norm, self.global_step) progress.set_description("Loss %.4f" % (loss_val)) if self.global_step % self.save_step == 0: self.save_model('states') mask_spec = self.up_sample_frames( spec_masked[0], return_first=True) pred_spec = self.up_sample_frames( pred_spec[0], return_first=True) true_spec = self.up_sample_frames( spec_stacked[0], return_first=True) mask_spec = plot_spectrogram_to_numpy( mask_spec.data.cpu().numpy()) pred_spec = plot_spectrogram_to_numpy( pred_spec.data.cpu().numpy()) true_spec = plot_spectrogram_to_numpy( true_spec.data.cpu().numpy()) self.log.add_image('mask_spec', mask_spec, self.global_step) self.log.add_image('pred_spec', pred_spec, self.global_step) self.log.add_image('true_spec', true_spec, self.global_step) loss_val = 0 pbar.update(1) self.global_step += 1 except RuntimeError as e: if 'CUDA out of memory' in str(e): print('CUDA out of memory at step: ', self.global_step) torch.cuda.empty_cache() self.optimizer.zero_grad() else: raise pbar.close() self.log.close()
def main(): parser = argparse.ArgumentParser() # Required parameters parser.add_argument('--pregenerated_data', type=Path, required=True) parser.add_argument('--teacher_model', default=None, type=str, required=True) parser.add_argument('--student_model', default=None, type=str, required=True) parser.add_argument('--output_dir', default=None, type=str, required=True) # Other parameters parser.add_argument( '--max_seq_length', default=128, type=int, help= 'The maximum total input sequence length after WordPiece tokenization. \n' 'Sequences longer than this will be truncated, and sequences shorter \n' 'than this will be padded.', ) parser.add_argument( '--reduce_memory', action='store_true', help= 'Store training data as on-disc memmaps to massively reduce memory usage', ) parser.add_argument( '--do_eval', action='store_true', help='Whether to run eval on the dev set.', ) parser.add_argument( '--do_lower_case', action='store_true', help='Set this flag if you are using an uncased model.', ) parser.add_argument( '--train_batch_size', default=32, type=int, help='Total batch size for training.', ) parser.add_argument( '--eval_batch_size', default=8, type=int, help='Total batch size for eval.', ) parser.add_argument( '--learning_rate', default=5e-5, type=float, help='The initial learning rate for Adam.', ) parser.add_argument( '--weight_decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay', ) parser.add_argument( '--num_train_epochs', default=3.0, type=float, help='Total number of training epochs to perform.', ) parser.add_argument( '--warmup_proportion', default=0.1, type=float, help= 'Proportion of training to perform linear learning rate warmup for. ' 'E.g., 0.1 = 10%% of training.', ) parser.add_argument( '--no_cuda', action='store_true', help='Whether not to use CUDA when available', ) parser.add_argument( '--local_rank', type=int, default=-1, help='local_rank for distributed training on gpus', ) parser.add_argument( '--seed', type=int, default=42, help='random seed for initialization', ) parser.add_argument( '--gradient_accumulation_steps', type=int, default=1, help= 'Number of updates steps to accumulate before performing a backward/update pass.', ) parser.add_argument( '--fp16', action='store_true', help='Whether to use 16-bit float precision instead of 32-bit', ) parser.add_argument( '--continue_train', action='store_true', help='Whether to train from checkpoints', ) # Additional arguments parser.add_argument('--eval_step', type=int, default=1000) # This is used for running on Huawei Cloud. parser.add_argument('--data_url', type=str, default='') args = parser.parse_args() logger.info('args:{}'.format(args)) samples_per_epoch = [] for i in range(int(args.num_train_epochs)): epoch_file = args.pregenerated_data / 'epoch_{}.json'.format(i) metrics_file = args.pregenerated_data / 'epoch_{}_metrics.json'.format( i) if epoch_file.is_file() and metrics_file.is_file(): metrics = json.loads(metrics_file.read_text()) samples_per_epoch.append(metrics['num_training_examples']) else: if i == 0: exit('No training data was found!') print( 'Warning! There are fewer epochs of pregenerated data ({}) than training epochs ({}).' .format(i, args.num_train_epochs)) print( 'This script will loop over the available data, but training diversity may be negatively impacted.' ) num_data_epochs = i break else: num_data_epochs = args.num_train_epochs if args.local_rank == -1 or args.no_cuda: device = torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu') n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device('cuda', args.local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logging.basicConfig( format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN, ) logger.info( 'device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}'. format(device, n_gpu, bool(args.local_rank != -1), args.fp16)) if args.gradient_accumulation_steps < 1: raise ValueError( 'Invalid gradient_accumulation_steps parameter: {}, should be >= 1' .format(args.gradient_accumulation_steps)) args.train_batch_size = (args.train_batch_size // args.gradient_accumulation_steps) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError( 'Output directory ({}) already exists and is not empty.'.format( args.output_dir)) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) total_train_examples = 0 for i in range(int(args.num_train_epochs)): # The modulo takes into account the fact that we may loop over limited epochs of data total_train_examples += samples_per_epoch[i % len(samples_per_epoch)] num_train_optimization_steps = int(total_train_examples / args.train_batch_size / args.gradient_accumulation_steps) if args.local_rank != -1: num_train_optimization_steps = (num_train_optimization_steps // torch.distributed.get_world_size()) if args.continue_train: student_model = TinyBertForPreTraining.from_pretrained( args.student_model) else: student_model = TinyBertForPreTraining.from_scratch(args.student_model) teacher_model = BertModel.from_pretrained(args.teacher_model) # student_model = TinyBertForPreTraining.from_scratch(args.student_model, fit_size=teacher_model.config.hidden_size) student_model.to(device) teacher_model.to(device) if args.local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError( 'Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.' ) teacher_model = DDP(teacher_model) elif n_gpu > 1: student_model = torch.nn.DataParallel(student_model) teacher_model = torch.nn.DataParallel(teacher_model) size = 0 for n, p in student_model.named_parameters(): logger.info('n: {}'.format(n)) logger.info('p: {}'.format(p.nelement())) size += p.nelement() logger.info('Total parameters: {}'.format(size)) # Prepare optimizer param_optimizer = list(student_model.named_parameters()) no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [ { 'params': [ p for n, p in param_optimizer if not any(nd in n for nd in no_decay) ], 'weight_decay': 0.01, }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0, }, ] loss_mse = MSELoss() optimizer = BertAdam( optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=num_train_optimization_steps, ) global_step = 0 logging.info('***** Running training *****') logging.info(' Num examples = {}'.format(total_train_examples)) logging.info(' Batch size = %d', args.train_batch_size) logging.info(' Num steps = %d', num_train_optimization_steps) for epoch in trange(int(args.num_train_epochs), desc='Epoch'): epoch_dataset = PregeneratedDataset( epoch=epoch, training_path=args.pregenerated_data, tokenizer=tokenizer, num_data_epochs=num_data_epochs, reduce_memory=args.reduce_memory, ) if args.local_rank == -1: train_sampler = RandomSampler(epoch_dataset) else: train_sampler = DistributedSampler(epoch_dataset) train_dataloader = DataLoader( epoch_dataset, sampler=train_sampler, batch_size=args.train_batch_size, ) tr_loss = 0.0 tr_att_loss = 0.0 tr_rep_loss = 0.0 student_model.train() nb_tr_examples, nb_tr_steps = 0, 0 with tqdm(total=len(train_dataloader), desc='Epoch {}'.format(epoch)) as pbar: for step, batch in enumerate( tqdm(train_dataloader, desc='Iteration', ascii=True)): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, lm_label_ids, is_next = ( batch) if input_ids.size()[0] != args.train_batch_size: continue att_loss = 0.0 rep_loss = 0.0 student_atts, student_reps = student_model( input_ids, segment_ids, input_mask) teacher_reps, teacher_atts, _ = teacher_model( input_ids, segment_ids, input_mask) teacher_reps = [ teacher_rep.detach() for teacher_rep in teacher_reps ] # speedup 1.5x teacher_atts = [ teacher_att.detach() for teacher_att in teacher_atts ] teacher_layer_num = len(teacher_atts) student_layer_num = len(student_atts) assert teacher_layer_num % student_layer_num == 0 layers_per_block = int(teacher_layer_num / student_layer_num) new_teacher_atts = [ teacher_atts[i * layers_per_block + layers_per_block - 1] for i in range(student_layer_num) ] for student_att, teacher_att in zip(student_atts, new_teacher_atts): student_att = torch.where( student_att <= -1e2, torch.zeros_like(student_att).to(device), student_att, ) teacher_att = torch.where( teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device), teacher_att, ) att_loss += loss_mse(student_att, teacher_att) new_teacher_reps = [ teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1) ] new_student_reps = student_reps for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps): rep_loss += loss_mse(student_rep, teacher_rep) loss = att_loss + rep_loss if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: optimizer.backward(loss) else: loss.backward() tr_att_loss += att_loss.item() tr_rep_loss += rep_loss.item() tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 pbar.update(1) mean_loss = (tr_loss * args.gradient_accumulation_steps / nb_tr_steps) mean_att_loss = (tr_att_loss * args.gradient_accumulation_steps / nb_tr_steps) mean_rep_loss = (tr_rep_loss * args.gradient_accumulation_steps / nb_tr_steps) if (step + 1) % args.gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() global_step += 1 if (global_step + 1) % args.eval_step == 0: result = {} result['global_step'] = global_step result['loss'] = mean_loss result['att_loss'] = mean_att_loss result['rep_loss'] = mean_rep_loss output_eval_file = os.path.join( args.output_dir, 'log.txt') with open(output_eval_file, 'a') as writer: logger.info('***** Eval results *****') for key in sorted(result.keys()): logger.info(' %s = %s', key, str(result[key])) writer.write('%s = %s\n' % (key, str(result[key]))) # Save a trained model model_name = 'step_{}_{}'.format( global_step, WEIGHTS_NAME) logging.info( '** ** * Saving fine-tuned model ** ** * ') # Only save the model it-self model_to_save = (student_model.module if hasattr( student_model, 'module') else student_model) output_model_file = os.path.join( args.output_dir, model_name) output_config_file = os.path.join( args.output_dir, CONFIG_NAME) torch.save(model_to_save.state_dict(), output_model_file) model_to_save.config.to_json_file(output_config_file) tokenizer.save_vocabulary(args.output_dir) if oncloud: logging.info( mox.file.list_directory(args.output_dir, recursive=True)) logging.info( mox.file.list_directory('.', recursive=True)) mox.file.copy_parallel(args.output_dir, args.data_url) mox.file.copy_parallel('.', args.data_url) model_name = 'step_{}_{}'.format(global_step, WEIGHTS_NAME) logging.info('** ** * Saving fine-tuned model ** ** * ') model_to_save = (student_model.module if hasattr( student_model, 'module') else student_model) output_model_file = os.path.join(args.output_dir, model_name) output_config_file = os.path.join(args.output_dir, CONFIG_NAME) torch.save(model_to_save.state_dict(), output_model_file) model_to_save.config.to_json_file(output_config_file) tokenizer.save_vocabulary(args.output_dir) if oncloud: logging.info( mox.file.list_directory(args.output_dir, recursive=True)) logging.info(mox.file.list_directory('.', recursive=True)) mox.file.copy_parallel(args.output_dir, args.data_url) mox.file.copy_parallel('.', args.data_url)
def main(): parser = argparse.ArgumentParser() parser.add_argument("--train_file_path", default=None, type=str, required=True) # Required parameters parser.add_argument("--teacher_model", default=None, type=str, required=True) parser.add_argument("--student_model", default=None, type=str, required=True) parser.add_argument("--output_dir", default=None, type=str, required=True) # Other parameters parser.add_argument( "--max_seq_len", default=128, type=int, help="The maximum total input sequence length after WordPiece \n" " tokenization. Sequences longer than this will be truncated, \n" "and sequences shorter than this will be padded.") parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.") parser.add_argument( "--do_lower_case", action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=32, type=int, help="Total batch size for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Total batch size for eval.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument('--weight_decay', '--wd', default=1e-1, type=float, metavar='W', help='weight decay') parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.") parser.add_argument( "--warmup_proportion", default=0.1, type=float, help= "Proportion of training to perform linear learning rate warmup for. " "E.g., 0.1 = 10%% of training.") parser.add_argument("--no_cuda", action='store_true', help="Whether not to use CUDA when available") parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument( '--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing \n" "a backward/update pass.") parser.add_argument( '--fp16', action='store_true', help="Whether to use 16-bit float precision instead of 32-bit") parser.add_argument('--continue_train', action='store_true', help='Whether to train from checkpoints') # Additional arguments parser.add_argument('--eval_step', type=int, default=1000) # This is used for running on Huawei Cloud. parser.add_argument('--data_url', type=str, default="") args = parser.parse_args() logger.info('args:{}'.format(args)) if args.local_rank == -1 or args.no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logging.basicConfig( format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN) logger.info( "device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}". format(device, n_gpu, bool(args.local_rank != -1), args.fp16)) if args.gradient_accumulation_steps < 1: raise ValueError( "Invalid gradient_accumulation_steps parameter: {}, should be >= 1" .format(args.gradient_accumulation_steps)) args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError( "Output directory ({}) already exists and is not empty.".format( args.output_dir)) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) tokenizer = BertTokenizer.from_pretrained(args.teacher_model, do_lower_case=args.do_lower_case) dataset = PregeneratedDataset(args.train_file_path, tokenizer, max_seq_len=args.max_seq_len) total_train_examples = len(dataset) num_train_optimization_steps = int( total_train_examples / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs) if args.local_rank != -1: num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size( ) * args.num_train_epochs if args.continue_train: student_model = TinyBertForPreTraining.from_pretrained( args.student_model) else: student_model = TinyBertForPreTraining.from_scratch(args.student_model) teacher_model = BertModel.from_pretrained(args.teacher_model) # student_model = TinyBertForPreTraining.from_scratch(args.student_model, fit_size=teacher_model.config.hidden_size) student_model.to(device) teacher_model.to(device) if args.local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training." ) teacher_model = DDP(teacher_model) elif n_gpu > 1: student_model = torch.nn.DataParallel(student_model) teacher_model = torch.nn.DataParallel(teacher_model) size = 0 for n, p in student_model.named_parameters(): logger.info('n: {}'.format(n)) logger.info('p: {}'.format(p.nelement())) size += p.nelement() logger.info('Total parameters: {}'.format(size)) # Prepare optimizer param_optimizer = list(student_model.named_parameters()) no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [{ 'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01 }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] loss_mse = MSELoss() optimizer = BertAdam(optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=num_train_optimization_steps) logging.info("***** Running training *****") logging.info(" Num examples = {}".format(total_train_examples)) logging.info(" Batch size = %d", args.train_batch_size) logging.info(" Num steps = %d", num_train_optimization_steps) if 1: if args.local_rank == -1: train_sampler = RandomSampler(dataset) else: train_sampler = DistributedSampler(dataset) train_dataloader = DataLoader(dataset, sampler=train_sampler, batch_size=args.train_batch_size) tr_loss = 0. tr_att_loss = 0. tr_rep_loss = 0. student_model.train() global_step = 0 nb_tr_examples, nb_tr_steps = 0, 0 for epoch in range(int(args.num_train_epochs)): for step, batch in enumerate( tqdm(train_dataloader, desc="Iteration", ascii=True)): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids = batch if input_ids.size()[0] != args.train_batch_size: continue att_loss = 0. rep_loss = 0. student_atts, student_reps = student_model( input_ids, segment_ids, input_mask) teacher_reps, teacher_atts, _ = teacher_model( input_ids, segment_ids, input_mask) # speedup 1.5x teacher_reps = [ teacher_rep.detach() for teacher_rep in teacher_reps ] teacher_atts = [ teacher_att.detach() for teacher_att in teacher_atts ] teacher_layer_num = len(teacher_atts) student_layer_num = len(student_atts) assert teacher_layer_num % student_layer_num == 0 layers_per_block = int(teacher_layer_num / student_layer_num) new_teacher_atts = [ teacher_atts[i * layers_per_block + layers_per_block - 1] for i in range(student_layer_num) ] for student_att, teacher_att in zip(student_atts, new_teacher_atts): student_att = torch.where( student_att <= -1e2, torch.zeros_like(student_att).to(device), student_att) teacher_att = torch.where( teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device), teacher_att) att_loss += loss_mse(student_att, teacher_att) new_teacher_reps = [ teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1) ] new_student_reps = student_reps for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps): rep_loss += loss_mse(student_rep, teacher_rep) loss = att_loss + rep_loss if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: optimizer.backward(loss) else: loss.backward() tr_att_loss += att_loss.item() tr_rep_loss += rep_loss.item() tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 mean_loss = tr_loss * args.gradient_accumulation_steps / nb_tr_steps mean_att_loss = tr_att_loss * args.gradient_accumulation_steps / nb_tr_steps mean_rep_loss = tr_rep_loss * args.gradient_accumulation_steps / nb_tr_steps if step % 100 == 0: logger.info(f'mean_loss = {mean_loss}') if (step + 1) % args.gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() global_step += 1 if (global_step + 1) % args.eval_step == 0: result = {} result['global_step'] = global_step result['loss'] = mean_loss result['att_loss'] = mean_att_loss result['rep_loss'] = mean_rep_loss output_eval_file = os.path.join( args.output_dir, "log.txt") with open(output_eval_file, "a") as writer: logger.info("***** Eval results *****") for key in sorted(result.keys()): logger.info(" %s = %s", key, str(result[key])) writer.write("%s = %s\n" % (key, str(result[key]))) # Save a trained model prefix = f"step_{step}" save_model(prefix, student_model, args.output_dir) prefix = f"epoch_{epoch}" save_model(prefix, student_model, args.output_dir)