(args, train_dataset, model, tokenizer)
return global_step, tr_loss / global_step
Train the model
Train the model
Train the model
def train(args, train_dataset, model, tokenizer):
    """ Train the model """
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriterP(args.output_dir)

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if p.requires_grad and not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if p.requires_grad and any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    warmup_steps = args.warmup_samples // args.train_batch_size
    if args.lr_decay:
        scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps=warmup_steps, t_total=t_total)
    else:
        scheduler = get_constant_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps)

    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d",
                   args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    try:
        with open(os.path.join(args.model_name_or_path, 'step.txt'), 'r') as c: 
            global_step = int(c.readline())
    except OSError as e:
        global_step = 0

    tr_loss, logging_loss = 0.0, 0.0
    moving_loss = MovingLoss(10000//args.logging_steps)
    model.zero_grad()

    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproducibility (even between python 2 and 3)
    try:    
        for _ in train_iterator:
            epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
            for step, batch in enumerate(epoch_iterator):
                inputs, labels = mask_tokens(batch, tokenizer, args) if args.mlm else (batch, batch)
                inputs = inputs.to(args.device)
                labels = labels.to(args.device)
                model.train()
                outputs = model(inputs, masked_lm_labels=labels) if args.mlm else model(inputs, labels=labels)
                loss = outputs[0]  # model outputs are always tuple in pytorch-transformers (see doc)

                if args.n_gpu > 1:
                    loss = loss.mean()  # mean() to average on multi-gpu parallel training
                if args.gradient_accumulation_steps > 1:
                    loss = loss / args.gradient_accumulation_steps

                if args.fp16:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                tr_loss += loss.item()
                moving_loss.add(loss.item())
                if (step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16:
                        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                    else:
                        torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                    optimizer.step()
                    scheduler.step()  # Update learning rate schedule
                    model.zero_grad()
                    global_step += 1

                    # Log metrics
                    if args.local_rank == -1 and args.evaluate_during_training and global_step % args.eval_steps == 0:  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, model, tokenizer, f"checkpoint-{global_step}")
                        for key, value in results.items():
                            tb_writer.add_scalar('eval_{}'.format(key), value, global_step)

                    if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                        tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
                        tb_writer.add_scalar('loss', (tr_loss - logging_loss)/args.logging_steps, global_step)
                        logging_loss = tr_loss
                        epoch_iterator.set_postfix(MovingLoss=f'{moving_loss.loss:.2f}', Perplexity=f'{torch.exp(torch.tensor(moving_loss.loss)):.2f}')

                    if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                        # Save model checkpoint
                        save_state(args, model, tokenizer, global_step)

                if args.max_steps > 0 and global_step > args.max_steps:
                    epoch_iterator.close()
                    break
            print_sample(model, tokenizer, args.device, args)
            if args.max_steps > 0 and global_step > args.max_steps:
                train_iterator.close()
                break
    except (KeyboardInterrupt, SystemExit):
        save_state(args, model, tokenizer, global_step)
        raise

    if args.local_rank in [-1, 0]:
        tb_writer.close()

    return global_step, tr_loss / global_step
def train ( args , train_dataset , model , tokenizer ) : if args . local_rank in [ - 1 , 0 ] : tb_writer = SummaryWriterP ( args . output_dir ) args . train_batch_size = args . per_gpu_train_batch_size * max ( 1 , args . n_gpu ) train_sampler = RandomSampler ( train_dataset ) if args . local_rank == - 1 else DistributedSampler ( train_dataset ) train_dataloader = DataLoader ( train_dataset , sampler = train_sampler , batch_size = args . train_batch_size ) if args . max_steps > 0 : t_total = args . max_steps args . num_train_epochs = args . max_steps // ( len ( train_dataloader ) // args . gradient_accumulation_steps ) + 1 else : t_total = len ( train_dataloader ) // args . gradient_accumulation_steps * args . num_train_epochs # Prepare optimizer and schedule (linear warmup and decay) no_decay = [ 'bias' , 'LayerNorm.weight' ] optimizer_grouped_parameters = [ { 'params' : [ p for n , p in model . named_parameters ( ) if p . requires_grad and not any ( nd in n for nd in no_decay ) ] , 'weight_decay' : args . weight_decay } , { 'params' : [ p for n , p in model . named_parameters ( ) if p . requires_grad and any ( nd in n for nd in no_decay ) ] , 'weight_decay' : 0.0 } ] optimizer = AdamW ( optimizer_grouped_parameters , lr = args . learning_rate , eps = args . adam_epsilon ) warmup_steps = args . warmup_samples // args . train_batch_size if args . lr_decay : scheduler = get_cosine_schedule_with_warmup ( optimizer , warmup_steps = warmup_steps , t_total = t_total ) else : scheduler = get_constant_schedule_with_warmup ( optimizer , num_warmup_steps = warmup_steps ) if args . fp16 : try : from apex import amp except ImportError : raise ImportError ( "Please install apex from https://www.github.com/nvidia/apex to use fp16 training." ) model , optimizer = amp . initialize ( model , optimizer , opt_level = args . fp16_opt_level ) # multi-gpu training (should be after apex fp16 initialization) if args . n_gpu > 1 : model = torch . nn . DataParallel ( model ) # Distributed training (should be after apex fp16 initialization) if args . local_rank != - 1 : model = torch . nn . parallel . DistributedDataParallel ( model , device_ids = [ args . local_rank ] , output_device = args . local_rank , find_unused_parameters = True ) # Train! logger . info ( "***** Running training *****" ) logger . info ( "  Num examples = %d" , len ( train_dataset ) ) logger . info ( "  Num Epochs = %d" , args . num_train_epochs ) logger . info ( "  Instantaneous batch size per GPU = %d" , args . per_gpu_train_batch_size ) logger . info ( "  Total train batch size (w. parallel, distributed & accumulation) = %d" , args . train_batch_size * args . gradient_accumulation_steps * ( torch . distributed . get_world_size ( ) if args . local_rank != - 1 else 1 ) ) logger . info ( "  Gradient Accumulation steps = %d" , args . gradient_accumulation_steps ) logger . info ( "  Total optimization steps = %d" , t_total ) try : with open ( os . path . join ( args . model_name_or_path , 'step.txt' ) , 'r' ) as c : global_step = int ( c . readline ( ) ) except OSError as e : global_step = 0 tr_loss , logging_loss = 0.0 , 0.0 moving_loss = MovingLoss ( 10000 // args . logging_steps ) model . zero_grad ( ) train_iterator = trange ( int ( args . num_train_epochs ) , desc = "Epoch" , disable = args . local_rank not in [ - 1 , 0 ] ) set_seed ( args ) # Added here for reproducibility (even between python 2 and 3) try : for _ in train_iterator : epoch_iterator = tqdm ( train_dataloader , desc = "Iteration" , disable = args . local_rank not in [ - 1 , 0 ] ) for step , batch in enumerate ( epoch_iterator ) : inputs , labels = mask_tokens ( batch , tokenizer , args ) if args . mlm else ( batch , batch ) inputs = inputs . to ( args . device ) labels = labels . to ( args . device ) model . train ( ) outputs = model ( inputs , masked_lm_labels = labels ) if args . mlm else model ( inputs , labels = labels ) loss = outputs [ 0 ] # model outputs are always tuple in pytorch-transformers (see doc) if args . n_gpu > 1 : loss = loss . mean ( ) # mean() to average on multi-gpu parallel training if args . gradient_accumulation_steps > 1 : loss = loss / args . gradient_accumulation_steps if args . fp16 : with amp . scale_loss ( loss , optimizer ) as scaled_loss : scaled_loss . backward ( ) else : loss . backward ( ) tr_loss += loss . item ( ) moving_loss . add ( loss . item ( ) ) if ( step + 1 ) % args . gradient_accumulation_steps == 0 : if args . fp16 : torch . nn . utils . clip_grad_norm_ ( amp . master_params ( optimizer ) , args . max_grad_norm ) else : torch . nn . utils . clip_grad_norm_ ( model . parameters ( ) , args . max_grad_norm ) optimizer . step ( ) scheduler . step ( ) # Update learning rate schedule model . zero_grad ( ) global_step += 1 # Log metrics if args . local_rank == - 1 and args . evaluate_during_training and global_step % args . eval_steps == 0 : # Only evaluate when single GPU otherwise metrics may not average well results = evaluate ( args , model , tokenizer , f"checkpoint-{global_step}" ) for key , value in results . items ( ) : tb_writer . add_scalar ( 'eval_{}' . format ( key ) , value , global_step ) if args . local_rank in [ - 1 , 0 ] and args . logging_steps > 0 and global_step % args . logging_steps == 0 : tb_writer . add_scalar ( 'lr' , scheduler . get_lr ( ) [ 0 ] , global_step ) tb_writer . add_scalar ( 'loss' , ( tr_loss - logging_loss ) / args . logging_steps , global_step ) logging_loss = tr_loss epoch_iterator . set_postfix ( MovingLoss = f'{moving_loss.loss:.2f}' , Perplexity = f'{torch.exp(torch.tensor(moving_loss.loss)):.2f}' ) if args . local_rank in [ - 1 , 0 ] and args . save_steps > 0 and global_step % args . save_steps == 0 : # Save model checkpoint save_state ( args , model , tokenizer , global_step ) if args . max_steps > 0 and global_step > args . max_steps : epoch_iterator . close ( ) break print_sample ( model , tokenizer , args . device , args ) if args . max_steps > 0 and global_step > args . max_steps : train_iterator . close ( ) break except ( KeyboardInterrupt , SystemExit ) : save_state ( args , model , tokenizer , global_step ) raise if args . local_rank in [ - 1 , 0 ] : tb_writer . close ( ) return global_step , tr_loss / global_step