def test_checkpointer_init(self): with tempfile.TemporaryDirectory() as temp_dir: more_config = { "checkpointing": True, "checkpointer_config": { "checkpoint_dir": temp_dir }, "log_writer_config": { "log_dir": temp_dir }, } trainer = Trainer(**base_config, **more_config, logging=True) trainer.fit(model, [dataloaders[0]]) self.assertIsNotNone(trainer.checkpointer) broken_config = { "checkpointing": True, "checkpointer_config": { "checkpoint_dir": None }, "log_writer_config": { "log_dir": temp_dir }, } with self.assertRaises(TypeError): trainer = Trainer(**base_config, **broken_config, logging=False) trainer.fit(model, [dataloaders[0]])
def test_log_writer_init(self): with tempfile.TemporaryDirectory() as temp_dir: log_writer_config = {"log_dir": temp_dir} trainer = Trainer( **base_config, logging=True, log_writer="json", log_writer_config=log_writer_config, ) trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.log_writer, LogWriter) log_writer_config = {"log_dir": temp_dir} trainer = Trainer( **base_config, logging=True, log_writer="tensorboard", log_writer_config=log_writer_config, ) trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.log_writer, TensorBoardWriter) log_writer_config = {"log_dir": temp_dir} with self.assertRaisesRegex(ValueError, "Unrecognized writer"): trainer = Trainer( **base_config, logging=True, log_writer="foo", log_writer_config=log_writer_config, ) trainer.fit(model, [dataloaders[0]])
def test_save_load(self): non_base_config = {"n_epochs": 2, "progress_bar": False} trainer1 = Trainer(**base_config, lr_scheduler="exponential") trainer1.fit(model, [dataloaders[0]]) trainer2 = Trainer(**non_base_config, lr_scheduler="linear") trainer3 = Trainer(**non_base_config, lr_scheduler="linear") with tempfile.NamedTemporaryFile() as fd: checkpoint_path = fd.name trainer1.save(checkpoint_path) trainer2.load(checkpoint_path, model=model) trainer3.load(checkpoint_path, None) self.assertEqual(trainer1.config, trainer2.config) self.dict_check( trainer1.optimizer.state_dict(), trainer2.optimizer.state_dict() ) # continue training after load trainer2.fit(model, [dataloaders[0]]) # check that an inappropriate model does not load an optimizer state but a trainer config self.assertEqual(trainer1.config, trainer3.config) self.assertFalse(hasattr(trainer3, "optimizer")) trainer3.fit(model, [dataloaders[0]])
def test_trainer_errors(self): dataloader = copy.deepcopy(dataloaders[0]) # No train split trainer = Trainer(**base_config) dataloader.dataset.split = "valid" with self.assertRaisesRegex(ValueError, "Cannot find any dataloaders"): trainer.fit(model, [dataloader]) # Unused split trainer = Trainer(**base_config, valid_split="val") with self.assertRaisesRegex(ValueError, "Dataloader splits must be"): trainer.fit(model, [dataloader])
def test_warmup(self): lr_scheduler_config = {"warmup_steps": 1, "warmup_unit": "batches"} trainer = Trainer(**base_config, lr_scheduler_config=lr_scheduler_config) trainer.fit(model, [dataloaders[0]]) self.assertEqual(trainer.warmup_steps, 1) lr_scheduler_config = {"warmup_steps": 1, "warmup_unit": "epochs"} trainer = Trainer(**base_config, lr_scheduler_config=lr_scheduler_config) trainer.fit(model, [dataloaders[0]]) self.assertEqual(trainer.warmup_steps, BATCHES_PER_EPOCH) lr_scheduler_config = {"warmup_percentage": 1 / BATCHES_PER_EPOCH} trainer = Trainer(**base_config, lr_scheduler_config=lr_scheduler_config) trainer.fit(model, [dataloaders[0]]) self.assertEqual(trainer.warmup_steps, 1)
def test_optimizer_init(self): trainer = Trainer(**base_config, optimizer="sgd") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.optimizer, optim.SGD) trainer = Trainer(**base_config, optimizer="adam") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.optimizer, optim.Adam) trainer = Trainer(**base_config, optimizer="adamax") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.optimizer, optim.Adamax) with self.assertRaisesRegex(ValueError, "Unrecognized optimizer"): trainer = Trainer(**base_config, optimizer="foo") trainer.fit(model, [dataloaders[0]])
def test_scheduler_init(self): trainer = Trainer(**base_config, lr_scheduler="constant") trainer.fit(model, [dataloaders[0]]) self.assertIsNone(trainer.lr_scheduler) trainer = Trainer(**base_config, lr_scheduler="linear") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.lr_scheduler, optim.lr_scheduler.LambdaLR) trainer = Trainer(**base_config, lr_scheduler="exponential") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.lr_scheduler, optim.lr_scheduler.ExponentialLR) trainer = Trainer(**base_config, lr_scheduler="step") trainer.fit(model, [dataloaders[0]]) self.assertIsInstance(trainer.lr_scheduler, optim.lr_scheduler.StepLR) with self.assertRaisesRegex(ValueError, "Unrecognized lr scheduler"): trainer = Trainer(**base_config, lr_scheduler="foo") trainer.fit(model, [dataloaders[0]])
def test_performance(self): """Test slicing performance with 2 corresponding slice tasks that represent roughly <10% of the data.""" dataloaders = [] for df, split in [(self.df_train, "train"), (self.df_valid, "valid")]: dataloader = create_dataloader(df, split) dataloaders.append(dataloader) base_task = create_task("task", module_suffixes=["A", "B"]) # Apply SFs slicing_functions = [f, g] # low-coverage slices slice_names = [sf.name for sf in slicing_functions] applier = PandasSFApplier(slicing_functions) S_train = applier.apply(self.df_train, progress_bar=False) S_valid = applier.apply(self.df_valid, progress_bar=False) # Add slice labels add_slice_labels(dataloaders[0], base_task, S_train) add_slice_labels(dataloaders[1], base_task, S_valid) # Convert to slice tasks tasks = convert_to_slice_tasks(base_task, slice_names) model = MultitaskClassifier(tasks=tasks) # Train # NOTE: Needs more epochs to convergence with more heads trainer = Trainer(lr=0.001, n_epochs=65, progress_bar=False) trainer.fit(model, dataloaders) scores = model.score(dataloaders) # Confirm reasonably high slice scores # Check train scores self.assertGreater(scores["task/TestData/train/f1"], 0.9) self.assertGreater(scores["task_slice:f_pred/TestData/train/f1"], 0.9) self.assertGreater(scores["task_slice:f_ind/TestData/train/f1"], 0.9) self.assertGreater(scores["task_slice:g_pred/TestData/train/f1"], 0.9) self.assertGreater(scores["task_slice:g_ind/TestData/train/f1"], 0.9) self.assertGreater(scores["task_slice:base_pred/TestData/train/f1"], 0.9) self.assertEqual(scores["task_slice:base_ind/TestData/train/f1"], 1.0) # Check valid scores self.assertGreater(scores["task/TestData/valid/f1"], 0.9) self.assertGreater(scores["task_slice:f_pred/TestData/valid/f1"], 0.9) self.assertGreater(scores["task_slice:f_ind/TestData/valid/f1"], 0.9) self.assertGreater(scores["task_slice:g_pred/TestData/valid/f1"], 0.9) self.assertGreater(scores["task_slice:g_ind/TestData/valid/f1"], 0.9) self.assertGreater(scores["task_slice:base_pred/TestData/valid/f1"], 0.9) # base_ind is trivial: all labels are positive self.assertEqual(scores["task_slice:base_ind/TestData/valid/f1"], 1.0)
def test_convergence(self): """Test slicing convergence with 1 slice task that represents ~25% of the data.""" dataloaders = [] for df, split in [(self.df_train, "train"), (self.df_valid, "valid")]: dataloader = create_dataloader(df, split) dataloaders.append(dataloader) base_task = create_task("task", module_suffixes=["A", "B"]) # Apply SFs slicing_functions = [h] # high coverage slice slice_names = [sf.name for sf in slicing_functions] applier = PandasSFApplier(slicing_functions) S_train = applier.apply(self.df_train, progress_bar=False) S_valid = applier.apply(self.df_valid, progress_bar=False) self.assertEqual(S_train.shape, (self.N_TRAIN, )) self.assertEqual(S_valid.shape, (self.N_VALID, )) self.assertIn("h", S_train.dtype.names) # Add slice labels add_slice_labels(dataloaders[0], base_task, S_train) add_slice_labels(dataloaders[1], base_task, S_valid) # Convert to slice tasks tasks = convert_to_slice_tasks(base_task, slice_names) model = MultitaskClassifier(tasks=tasks) # Train trainer = Trainer(lr=0.001, n_epochs=50, progress_bar=False) trainer.fit(model, dataloaders) scores = model.score(dataloaders) # Confirm near perfect scores self.assertGreater(scores["task/TestData/valid/accuracy"], 0.94) self.assertGreater(scores["task_slice:h_pred/TestData/valid/accuracy"], 0.94) self.assertGreater(scores["task_slice:h_ind/TestData/valid/f1"], 0.94) # Calculate/check train/val loss train_dataset = dataloaders[0].dataset train_loss_output = model.calculate_loss(train_dataset.X_dict, train_dataset.Y_dict) train_loss = train_loss_output[0]["task"].item() self.assertLess(train_loss, 0.1) val_dataset = dataloaders[1].dataset val_loss_output = model.calculate_loss(val_dataset.X_dict, val_dataset.Y_dict) val_loss = val_loss_output[0]["task"].item() self.assertLess(val_loss, 0.1)
def test_log_writer_json(self): # Addresses issue #1439 # Confirm that a log file is written to the specified location after training run_name = "log.json" with tempfile.TemporaryDirectory() as temp_dir: log_writer_config = {"log_dir": temp_dir, "run_name": run_name} trainer = Trainer( **base_config, logging=True, log_writer="json", log_writer_config=log_writer_config, ) trainer.fit(model, [dataloaders[0]]) log_path = os.path.join(trainer.log_writer.log_dir, run_name) with open(log_path, "r") as f: log = json.load(f) self.assertIn("model/all/train/loss", log)
def test_convergence(self): """Test multitask classifier convergence with two tasks.""" dataloaders = [] for offset, task_name in zip([0.0, 0.25], ["task1", "task2"]): df = create_data(N_TRAIN, offset) dataloader = create_dataloader(df, "train", task_name) dataloaders.append(dataloader) for offset, task_name in zip([0.0, 0.25], ["task1", "task2"]): df = create_data(N_VALID, offset) dataloader = create_dataloader(df, "valid", task_name) dataloaders.append(dataloader) task1 = create_task("task1", module_suffixes=["A", "A"]) task2 = create_task("task2", module_suffixes=["A", "B"]) model = MultitaskClassifier(tasks=[task1, task2]) # Train trainer = Trainer(lr=0.001, n_epochs=10, progress_bar=False) trainer.fit(model, dataloaders) scores = model.score(dataloaders) # Confirm near perfect scores on both tasks for idx, task_name in enumerate(["task1", "task2"]): self.assertGreater(scores[f"{task_name}/TestData/valid/accuracy"], 0.95) # Calculate/check train/val loss train_dataset = dataloaders[idx].dataset train_loss_output = model.calculate_loss( train_dataset.X_dict, train_dataset.Y_dict ) train_loss = train_loss_output[0][task_name].item() self.assertLess(train_loss, 0.05) val_dataset = dataloaders[2 + idx].dataset val_loss_output = model.calculate_loss( val_dataset.X_dict, val_dataset.Y_dict ) val_loss = val_loss_output[0][task_name].item() self.assertLess(val_loss, 0.05)
test_dl_slice = slice_model.make_slice_dataloader(test_dl.dataset, S_test, shuffle=False, batch_size=BATCH_SIZE) # %% [markdown] # ### Representation learning with slices # %% [markdown] # Using Snorkel's [`Trainer`](https://snorkel.readthedocs.io/en/master/packages/_autosummary/classification/snorkel.classification.Trainer.html), we fit our classifier with the training set dataloader. # %% from snorkel.classification import Trainer # For demonstration purposes, we set n_epochs=2 trainer = Trainer(n_epochs=2, lr=1e-4, progress_bar=True) trainer.fit(slice_model, [train_dl_slice]) # %% [markdown] # At inference time, the primary task head (`spam_task`) will make all final predictions. # We'd like to evaluate all the slice heads on the original task head — [`score_slices`](https://snorkel.readthedocs.io/en/v0.9.3/packages/_autosummary/slicing/snorkel.slicing.SliceAwareClassifier.html#snorkel.slicing.SliceAwareClassifier.score_slices) remaps all slice-related labels, denoted `spam_task_slice:{slice_name}_pred`, to be evaluated on the `spam_task`. # %% slice_model.score_slices([test_dl_slice], as_dataframe=True) # %% [markdown] # *Note: in this toy dataset, we see high variance in slice performance, because our dataset is so small that (i) there are few data points in the train split, giving little signal to learn over, and (ii) there are few data points in the test split, making our evaluation metrics very noisy. # For a demonstration of data slicing deployed in state-of-the-art models, please see our [SuperGLUE](https://github.com/HazyResearch/snorkel-superglue/tree/master/tutorials) tutorials.* # %% [markdown] # ---
# %% import torchvision.models as models # initialize pretrained feature extractor cnn = models.resnet18(pretrained=True) model = create_model(cnn) # %% [markdown] # ### Train and Evaluate Model # %% {"tags": ["md-exclude-output"]} from snorkel.classification import Trainer trainer = Trainer( n_epochs=1, # increase for improved performance lr=1e-3, checkpointing=True, checkpointer_config={"checkpoint_dir": "checkpoint"}, ) trainer.fit(model, [dl_train]) # %% model.score([dl_valid]) # %% [markdown] # ## Recap # We have successfully trained a visual relationship detection model! Using categorical and spatial intuition about how objects in a visual relationship interact with each other, we are able to assign high quality training labels to object pairs in the VRD dataset in a multi-class classification setting. # # For more on how Snorkel can be used for visual relationship tasks, please see our [ICCV 2019 paper](https://arxiv.org/abs/1904.11622)!
def test_trainer_onetask(self): """Train a single-task model""" trainer = Trainer(**base_config) trainer.fit(model, [dataloaders[0]])
def slicing_evaluation(df_train, df_test, train_model=None): if train_model is None: train_model = "mlp" sfs = [ SlicingFunction.short_comment, SlicingFunction.ind_keyword, SlicingFunction.cmp_re, SlicingFunction.industry_keyword ] slice_names = [sf.name for sf in sfs] scorer = Scorer(metrics=["f1"]) ft = FT.load(f"{WORK_PATH}/snorkel_flow/sources/fasttext_name_model.bin") def get_ftr(text): return ft.get_sentence_vector(' '.join( [w for w in jieba.lcut(text.strip())])) X_train = np.array(list(df_train.text.apply(get_ftr).values)) X_test = np.array(list(df_test.text.apply(get_ftr).values)) Y_train = df_train.label.values Y_test = df_test.label.values if train_model == "lr": sklearn_model = LogisticRegression(C=0.001, solver="liblinear") sklearn_model.fit(X=X_train, y=Y_train) preds_test = sklearn_model.predict(X_test) probs_test = preds_to_probs( preds_test, len([c for c in dir(Polarity) if not c.startswith("__")])) print(f"Test set F1: {100 * f1_score(Y_test, preds_test):.1f}%") applier = PandasSFApplier(sfs) S_test = applier.apply(df_test) analysis = scorer.score_slices(S=S_test, golds=Y_test, preds=preds_test, probs=probs_test, as_dataframe=True) return analysis if train_model == "mlp": # Define model architecture bow_dim = X_train.shape[1] hidden_dim = bow_dim mlp = get_pytorch_mlp(hidden_dim=hidden_dim, num_layers=2) # Initialize slice model slice_model = SliceAwareClassifier( base_architecture=mlp, head_dim=hidden_dim, slice_names=slice_names, scorer=scorer, ) # generate the remaining S matrices with the new set of slicing functions applier = PandasSFApplier(sfs) S_train = applier.apply(df_train) S_test = applier.apply(df_test) # add slice labels to an existing dataloader BATCH_SIZE = 64 train_dl = create_dict_dataloader(X_train, Y_train, "train") train_dl_slice = slice_model.make_slice_dataloader( train_dl.dataset, S_train, shuffle=True, batch_size=BATCH_SIZE) test_dl = create_dict_dataloader(X_test, Y_test, "train") test_dl_slice = slice_model.make_slice_dataloader( test_dl.dataset, S_test, shuffle=False, batch_size=BATCH_SIZE) # fit our classifier with the training set dataloader trainer = Trainer(n_epochs=2, lr=1e-4, progress_bar=True) trainer.fit(slice_model, [train_dl_slice]) analysis = slice_model.score_slices([test_dl_slice], as_dataframe=True) return analysis
def test_trainer_twotask(self): """Train a model with overlapping modules and flows""" multitask_model = MultitaskClassifier(tasks) trainer = Trainer(**base_config) trainer.fit(multitask_model, dataloaders)
) # Add task to list of tasks tasks.append(task_object) # Input list of tasks to MultitaskClassifier object to create model with architecture set for each task model = MultitaskClassifier(tasks) # Set out trainer settings - I.e. how the model will train trainer_config = { "progress_bar": True, "n_epochs": 2, "lr": 0.02, "logging": True, "log_writer": "json", "checkpointing": True, } # Create trainer object using above settings trainer = Trainer(**trainer_config) # Train model using above settings on the datasets linked trainer.fit(model, dataloaders) # Output training stats of model trainer.log_writer.write_log("output_statistics.json") # Score model using test set and print model_scores = model.score(dataloaders) print(model_scores)
def train(args, train_dataset, model, tokenizer): """ Train the model """ if args.local_rank in [-1, 0]: tb_writer = SummaryWriter() 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 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 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) scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total) 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) global_step = 0 tr_loss, logging_loss = 0.0, 0.0 model.zero_grad() train_iterator = range(int(args.num_train_epochs)) set_seed( args) # Added here for reproductibility (even between python 2 and 3) if args.model_type == 'bert-slice-aware' or args.model_type == 'bert-slice-aware-random-slices': if args.model_type == 'bert-slice-aware': sfs = slicing_functions[args.task_name] elif args.model_type == 'bert-slice-aware-random-slices': if args.number_random_slices is None or args.size_random_slices is None: sfs = random_slicing_functions[args.task_name] else: sfs = args.sfs processor = slicing_processors[args.task_name]() examples_train = processor.get_train_examples(args.data_dir) snorkel_sf_applier = SFApplier(sfs) if os.path.isfile(args.data_dir + "/snorkel_slices_train.pickle"): with open(args.data_dir + "/snorkel_slices_train.pickle", "rb") as f: logger.info("loaded cached pickle for sliced train.") snorkel_slices_train = pickle.load(f) else: snorkel_slices_train = snorkel_sf_applier.apply(examples_train) with open(args.data_dir + "/snorkel_slices_train.pickle", "wb") as f: pickle.dump(snorkel_slices_train, f) logger.info("dumped pickle with sliced train.") snorkel_slices_with_ns = [] for i, example in enumerate(examples_train): for _ in range(len(example.documents)): snorkel_slices_with_ns.append(snorkel_slices_train[i]) snorkel_slices_with_ns_np = np.array(snorkel_slices_with_ns, dtype=snorkel_slices_train.dtype) slice_model = SliceAwareClassifier( task_name='labels', input_data_key='input_ids', base_architecture=model, head_dim=768, #* args.max_seq_length, slice_names=[sf.name for sf in sfs]) X_dict = { 'input_ids': train_dataset.tensors[0], 'attention_mask': train_dataset.tensors[1], 'token_type_ids': train_dataset.tensors[2] } Y_dict = {'labels': train_dataset.tensors[3]} ds = DictDataset(name='labels', split='train', X_dict=X_dict, Y_dict=Y_dict) train_dl_slice = slice_model.make_slice_dataloader( ds, snorkel_slices_with_ns_np, shuffle=True, batch_size=args.train_batch_size) trainer = Trainer(lr=args.learning_rate, n_epochs=int(args.num_train_epochs), l2=args.weight_decay, optimizer="adamax", max_steps=args.max_steps, seed=args.seed) trainer.fit(slice_model, [train_dl_slice]) model = slice_model else: for _ in train_iterator: epoch_iterator = train_dataloader for step, batch in enumerate(epoch_iterator): model.train() batch = tuple(t.to(args.device) for t in batch) inputs = { 'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3] } if args.model_type != 'distilbert': inputs['token_type_ids'] = batch[2] if args.model_type in [ 'bert', 'xlnet' ] else None # XLM, DistilBERT and RoBERTa don't use segment_ids if args.model_type == 'bert-mtl': inputs["clf_head"] = 0 outputs = model(**inputs) loss = outputs[ 0] # model outputs are always tuple in 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() torch.nn.utils.clip_grad_norm_( amp.master_params(optimizer), args.max_grad_norm) else: loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) tr_loss += loss.item() if (step + 1) % args.gradient_accumulation_steps == 0: optimizer.step() scheduler.step() # Update learning rate schedule model.zero_grad() global_step += 1 if args.local_rank in [ -1, 0 ] and args.logging_steps > 0 and global_step % args.logging_steps == 0: # Log metrics if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well results = evaluate(args, model, tokenizer, sample_percentage=0.01) for key, value in results.items(): tb_writer.add_scalar('eval_{}'.format(key), value, global_step) ex.log_scalar('eval_{}'.format(key), value, global_step) logger.info('eval_{}'.format(key) + ": " + str(value) + ", step: " + str(global_step)) 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) ex.log_scalar("lr", scheduler.get_lr()[0], global_step) ex.log_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step) logging_loss = tr_loss if args.local_rank in [ -1, 0 ] and args.save_steps > 0 and global_step % args.save_steps == 0: # Save model checkpoint output_dir = os.path.join( args.output_dir, 'checkpoint-{}'.format(global_step)) if not os.path.exists(output_dir): os.makedirs(output_dir) model_to_save = model.module if hasattr( model, 'module' ) else model # Take care of distributed/parallel training model_to_save.save_pretrained(output_dir) torch.save( args, os.path.join(output_dir, 'training_args.bin')) logger.info("Saving model checkpoint to %s", output_dir) if args.max_steps > 0 and global_step > args.max_steps: break # epoch_iterator.close() if args.max_steps > 0 and global_step > args.max_steps: break # train_iterator.close() if args.local_rank in [-1, 0]: tb_writer.close() return model