update to store uncompressed model on S3

examples/run_classifier.py

@@ -24,6 +24,7 @@ import os
 import logging
 import argparse
 import random
+from datetime import datetime
 from tqdm import tqdm, trange
 
 import numpy as np
@@ -580,6 +581,7 @@ def main():
         model.eval()
         eval_loss, eval_accuracy = 0, 0
         nb_eval_steps, nb_eval_examples = 0, 0
+        start_time = datetime.now()
         for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
             input_ids = input_ids.to(device)
             input_mask = input_mask.to(device)
@@ -600,6 +602,10 @@ def main():
             nb_eval_examples += input_ids.size(0)
             nb_eval_steps += 1
 
+        total_time = datetime.now() - start_time
+        logger.info("Total time: %d seconds — Number of examples: %d — Examples/second: %f" % (
+            total_time.total_seconds(), len(eval_data), len(eval_data)/total_time.total_seconds()))
+
         eval_loss = eval_loss / nb_eval_steps
         eval_accuracy = eval_accuracy / nb_eval_examples
 

examples/run_squad.py

@@ -27,6 +27,7 @@ import math
 import os
 import random
 import pickle
+from datetime import datetime
 from tqdm import tqdm, trange
 
 import numpy as np
@@ -699,6 +700,7 @@ def main():
                              "be truncated to this length.")
     parser.add_argument("--do_train", default=False, action='store_true', help="Whether to run training.")
     parser.add_argument("--do_predict", default=False, action='store_true', help="Whether to run eval on the dev set.")
+    parser.add_argument("--do_benchmark", default=False, action='store_true', help="Whether to benchmark prediction speed on the dev set.")
     parser.add_argument("--train_batch_size", default=32, type=int, help="Total batch size for training.")
     parser.add_argument("--predict_batch_size", default=8, type=int, help="Total batch size for predictions.")
     parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
@@ -740,6 +742,10 @@ def main():
                         default=False,
                         action='store_true',
                         help="Whether to use 16-bit float precision instead of 32-bit")
+    parser.add_argument('--predict_fp16',
+                        default=False,
+                        action='store_true',
+                        help="Whether to use 16-bit float precision instead of 32-bit for prediction")
     parser.add_argument('--loss_scale',
                         type=float, default=0,
                         help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
@@ -787,6 +793,7 @@ def main():
     if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
         raise ValueError("Output directory () already exists and is not empty.")
     os.makedirs(args.output_dir, exist_ok=True)
+    output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
 
     tokenizer = BertTokenizer.from_pretrained(args.bert_model)
 
@@ -799,6 +806,7 @@ def main():
             len(train_examples) / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs)
 
     # Prepare model
+    cache_folder
     model = BertForQuestionAnswering.from_pretrained(args.bert_model,
                 cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank))
 
@@ -915,17 +923,18 @@ def main():
                     optimizer.zero_grad()
                     global_step += 1
 
-    # Save a trained model
-    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
-    output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
-    torch.save(model_to_save.state_dict(), output_model_file)
+        # Save a trained model
+        model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
+        torch.save(model_to_save.state_dict(), output_model_file)
 
-    # Load a trained model that you have fine-tuned
-    model_state_dict = torch.load(output_model_file)
-    model = BertForQuestionAnswering.from_pretrained(args.bert_model, state_dict=model_state_dict)
-    model.to(device)
+    if (args.do_predict or args.do_benchmark) and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
+        # Load a trained model that you have fine-tuned
+        model_state_dict = torch.load(output_model_file)
+        model = BertForQuestionAnswering.from_pretrained(args.bert_model, state_dict=model_state_dict)
+        model.to(device)
+        if args.predict_fp16:
+            model.half()
 
-    if args.do_predict and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
         eval_examples = read_squad_examples(
             input_file=args.predict_file, is_training=False)
         eval_features = convert_examples_to_features(
@@ -953,28 +962,37 @@ def main():
         model.eval()
         all_results = []
         logger.info("Start evaluating")
+        start_time = datetime.now()
         for input_ids, input_mask, segment_ids, example_indices in tqdm(eval_dataloader, desc="Evaluating"):
             if len(all_results) % 1000 == 0:
                 logger.info("Processing example: %d" % (len(all_results)))
             input_ids = input_ids.to(device)
             input_mask = input_mask.to(device)
             segment_ids = segment_ids.to(device)
-            with torch.no_grad():
-                batch_start_logits, batch_end_logits = model(input_ids, segment_ids, input_mask)
-            for i, example_index in enumerate(example_indices):
-                start_logits = batch_start_logits[i].detach().cpu().tolist()
-                end_logits = batch_end_logits[i].detach().cpu().tolist()
-                eval_feature = eval_features[example_index.item()]
-                unique_id = int(eval_feature.unique_id)
-                all_results.append(RawResult(unique_id=unique_id,
-                                             start_logits=start_logits,
-                                             end_logits=end_logits))
-        output_prediction_file = os.path.join(args.output_dir, "predictions.json")
-        output_nbest_file = os.path.join(args.output_dir, "nbest_predictions.json")
-        write_predictions(eval_examples, eval_features, all_results,
-                          args.n_best_size, args.max_answer_length,
-                          args.do_lower_case, output_prediction_file,
-                          output_nbest_file, args.verbose_logging)
+            if args.do_benchmark:
+                with torch.no_grad():
+                    model(input_ids, segment_ids, input_mask)
+            else:
+                with torch.no_grad():
+                    batch_start_logits, batch_end_logits = model(input_ids, segment_ids, input_mask)
+                for i, example_index in enumerate(example_indices):
+                    start_logits = batch_start_logits[i].detach().cpu().tolist()
+                    end_logits = batch_end_logits[i].detach().cpu().tolist()
+                    eval_feature = eval_features[example_index.item()]
+                    unique_id = int(eval_feature.unique_id)
+                    all_results.append(RawResult(unique_id=unique_id,
+                                                start_logits=start_logits,
+                                                end_logits=end_logits))
+        total_time = datetime.now() - start_time
+        logger.info("Total time: %d seconds — Number of examples: %d — Examples/second: %f" % (
+            total_time.total_seconds(), len(eval_data), len(eval_data)/total_time.total_seconds()))
+        if not args.do_benchmark:
+            output_prediction_file = os.path.join(args.output_dir, "predictions.json")
+            output_nbest_file = os.path.join(args.output_dir, "nbest_predictions.json")
+            write_predictions(eval_examples, eval_features, all_results,
+                            args.n_best_size, args.max_answer_length,
+                            args.do_lower_case, output_prediction_file,
+                            output_nbest_file, args.verbose_logging)
 
 
 if __name__ == "__main__":

pytorch_pretrained_bert/modeling.py

@@ -36,15 +36,26 @@ from .file_utils import cached_path
 
 logger = logging.getLogger(__name__)
 
-PRETRAINED_MODEL_ARCHIVE_MAP = {
-    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased.tar.gz",
-    'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased.tar.gz",
-    'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased.tar.gz",
-    'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased.tar.gz",
-    'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased.tar.gz",
-    'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased.tar.gz",
-    'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese.tar.gz",
+PRETRAINED_WEIGHTS_ARCHIVE_MAP = {
+    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch-model.bin",
+    'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch-model.bin",
+    'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch-model.bin",
+    'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch-model.bin",
+    'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch-model.bin",
+    'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch-model.bin",
+    'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch-model.bin",
 }
+
+PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
+    'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json",
+    'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json",
+    'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json",
+    'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json",
+    'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json",
+    'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json",
+}
+
 CONFIG_NAME = 'bert_config.json'
 WEIGHTS_NAME = 'pytorch_model.bin'
 
@@ -466,47 +477,51 @@ class PreTrainedBertModel(nn.Module):
             *inputs, **kwargs: additional input for the specific Bert class
                 (ex: num_labels for BertForSequenceClassification)
         """
-        if pretrained_model_name in PRETRAINED_MODEL_ARCHIVE_MAP:
-            archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name]
-        else:
-            archive_file = pretrained_model_name
+        # Load config
+        config_file = PRETRAINED_CONFIG_ARCHIVE_MAP.get(pretrained_model_name, pretrained_model_name)
+        if os.path.isdir(config_file):
+            config_file = os.path.join(config_file, CONFIG_NAME)
         # redirect to the cache, if necessary
         try:
-            resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
+            resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
         except FileNotFoundError:
             logger.error(
                 "Model name '{}' was not found in model name list ({}). "
                 "We assumed '{}' was a path or url but couldn't find any file "
                 "associated to this path or url.".format(
                     pretrained_model_name,
-                    ', '.join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()),
-                    archive_file))
+                    ', '.join(PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
+                    config_file))
             return None
-        if resolved_archive_file == archive_file:
-            logger.info("loading archive file {}".format(archive_file))
-        else:
-            logger.info("loading archive file {} from cache at {}".format(
-                archive_file, resolved_archive_file))
-        tempdir = None
-        if os.path.isdir(resolved_archive_file):
-            serialization_dir = resolved_archive_file
-        else:
-            # Extract archive to temp dir
-            tempdir = tempfile.mkdtemp()
-            logger.info("extracting archive file {} to temp dir {}".format(
-                resolved_archive_file, tempdir))
-            with tarfile.open(resolved_archive_file, 'r:gz') as archive:
-                archive.extractall(tempdir)
-            serialization_dir = tempdir
-        # Load config
-        config_file = os.path.join(serialization_dir, CONFIG_NAME)
+        logger.info("loading config file {}{}".format(
+            "from cache " if resolved_config_file != config_file else " ", resolved_config_file
+        ))
         config = BertConfig.from_json_file(config_file)
         logger.info("Model config {}".format(config))
+
         # Instantiate model.
         model = cls(config, *inputs, **kwargs)
         if state_dict is None:
-            weights_path = os.path.join(serialization_dir, WEIGHTS_NAME)
-            state_dict = torch.load(weights_path)
+            weights_file = PRETRAINED_WEIGHTS_ARCHIVE_MAP.get(pretrained_model_name, pretrained_model_name)
+            if os.path.isdir(weights_file):
+                weights_file = os.path.join(weights_file, WEIGHTS_NAME)
+            try:
+                resolved_weights_file = cached_path(weights_file, cache_dir=cache_dir)
+            except FileNotFoundError:
+                logger.error(
+                    "Model name '{}' was not found in model name list ({}). "
+                    "We assumed '{}' was a path or url but couldn't find any file "
+                    "associated to this path or url.".format(
+                        pretrained_model_name,
+                        ', '.join(PRETRAINED_WEIGHTS_ARCHIVE_MAP.keys()),
+                        weights_file))
+                return None
+            logger.info("loading pretrained weights {}{}".format(
+                "from cache " if resolved_weights_file != weights_file else " ", resolved_weights_file
+            ))
+            if os.path.isdir(weights_file):
+                weights_file = os.path.join(weights_file, WEIGHTS_NAME)
+            state_dict = torch.load(weights_file)
 
         old_keys = []
         new_keys = []
@@ -545,9 +560,6 @@ class PreTrainedBertModel(nn.Module):
         if len(unexpected_keys) > 0:
             logger.info("Weights from pretrained model not used in {}: {}".format(
                 model.__class__.__name__, unexpected_keys))
-        if tempdir:
-            # Clean up temp dir
-            shutil.rmtree(tempdir)
         return model