Merge pull request haoyuhu#20 from HaoyuHu/pr17

haoyuhu · web-flow · commit e29e77d56c0e · 2019-12-13T19:06:15.000+08:00
PR#17
diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@ Feel free to fine tune large BERT models with large batch size easily. Multi-GPU
 
 - Tensorflow
   - tensorflow >= 1.11.0   # CPU Version of TensorFlow.
-  - tensorflow-gpu  >= 1.11.0  # GPU version of TensorFlow.
+  - tensorflow-gpu  >= 1.11.0  # GPU version of TensorFlow. (Upgrade to 1.14.0 when meets [ImportError: No module named 'tensorflow.python.distribute.cross_device_ops' ](https://github.com/HaoyuHu/bert-multi-gpu/issues/11))
 - NVIDIA Collective Communications Library (NCCL)
 
 
@@ -15,6 +15,14 @@ Feel free to fine tune large BERT models with large batch size easily. Multi-GPU
 
 - CPU/GPU/TPU Support
 - **Multi-GPU Support**: [`tf.distribute.MirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy) is used to achieve Multi-GPU support for this project, which mirrors vars to distribute across multiple devices and machines. The maximum batch_size for each GPU is almost the same as [bert](https://github.com/google-research/bert/blob/master/README.md#out-of-memory-issues). So **global batch_size** depends on how many GPUs there are.
+    - Assume: num_train_examples = 32000
+    - Situation 1 (multi-gpu): train_batch_size = 8, num_gpu_cores = 4, num_train_epochs = 1
+        - global_batch_size = train_batch_size * num_gpu_cores = 32
+        - iteration_steps = num_train_examples * num_train_epochs / train_batch_size = 4000
+    - Situation 2 (single-gpu): train_batch_size = 32, num_gpu_cores = 1, num_train_epochs = 4
+        - global_batch_size = train_batch_size * num_gpu_cores = 32
+        - iteration_steps = num_train_examples * num_train_epochs / train_batch_size = 4000
+    - Result after training is equivalent between situation 1 and 2 when synchronous update on gradients is applied.
 - **FP16 Support**: [FP16](https://en.wikipedia.org/wiki/Half-precision_floating-point_format) allows you to use a larger batch_size. And training speed will increase by 70~100% on Volta GPUs, but may be slower on Pascal GPUs([REF1](https://github.com/tensorflow/tensorflow/issues/15585#issuecomment-361769151), [REF2](https://github.com/HaoyuHu/bert-multi-gpu/issues/1#issuecomment-493363383)).
 - **SavedModel Export**
 
@@ -61,11 +69,28 @@ python run_custom_classifier.py \
   --num_train_epochs=3.0 \
   --use_gpu=true \
   --num_gpu_cores=4 \
-  --use_fp16=true \
+  --use_fp16=false \
   --output_dir=/cfs/outputs/bert-large-uncased-qqp
 ```
 
-
+Shell script is available also (see run_custom_classifier.sh)
+- Optional params could be passed flexibly through command line.
+- CUDA_VISIBLE_DEVICES could be set and export as environmental variables when multi-gpus are used.
+```shell
+# refer to the variables acronym
+bash run_custom_classifier.sh -h
+# output
+current params setting:
+-s max_seq_length,        default val is: 128
+-g num_gpu_cores,         default val is: 4
+-b train_batch_size,      default val is: 32
+-l learning_rate,         default val is: 2e-5
+-e num_train_epochs,      default val is: 3.0
+-c CUDA_VISIBLE_DEVICES,  default val is: 0,1,2,3
+
+# example to pass params
+bash run_custom_classifier.sh -s 512 -b 8 -l 3e-5 -e 1 -g 2 -c 2,3
+```
 
 ### Run Sequence Labeling
 
@@ -106,7 +131,7 @@ python run_seq_labeling.py \
   --num_train_epochs=10.0 \
   --use_gpu=true \
   --num_gpu_cores=4 \
-  --use_fp16=true \
+  --use_fp16=false \
   --output_dir=/cfs/outputs/bert-large-uncased-punct
 ```
 
diff --git a/run_custom_classifier.py b/run_custom_classifier.py
@@ -20,15 +20,19 @@
 
 import collections
 import csv
+import json
 import os
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.distribute.cross_device_ops import AllReduceCrossDeviceOps
+from tensorflow.python.estimator.estimator import Estimator
+from tensorflow.python.estimator.run_config import RunConfig
+
+import custom_optimization
 import modeling
 import optimization
-import custom_optimization
 import tokenization
-from tensorflow.python.distribute.cross_device_ops import AllReduceCrossDeviceOps
-import tensorflow as tf
-from tensorflow.python.estimator.run_config import RunConfig
-from tensorflow.python.estimator.estimator import Estimator
 
 flags = tf.flags
 
@@ -749,35 +753,35 @@ def metric_fn(per_example_loss, label_ids, logits, is_real_example):
                 predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
                 accuracy = tf.metrics.accuracy(
                     labels=label_ids, predictions=predictions, weights=is_real_example)
+                # add more metrics
+                pr, pr_op = tf.metrics.precision(
+                    labels=label_ids, predictions=predictions, weights=is_real_example)
+                re, re_op = tf.metrics.recall(
+                    labels=label_ids, predictions=predictions, weights=is_real_example)
+                f1 = (2 * pr * re) / (pr + re)  # f1-score for binary classification
                 loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
                 return {
                     "eval_accuracy": accuracy,
-                    "eval_loss": loss,
+                    "eval_precision": (pr, pr_op),
+                    "eval_recall": (re, re_op),
+                    "eval_f1score": (f1, tf.identity(f1)),
+                    "eval_loss": loss
                 }
 
             eval_metrics = (metric_fn,
                             [per_example_loss, label_ids, logits, is_real_example])
-            if use_gpu and int(num_gpu_cores) >= 2:
-                output_spec = tf.estimator.EstimatorSpec(
-                    mode=mode,
-                    loss=total_loss,
-                    eval_metric_ops=eval_metrics[0](*eval_metrics[1]))
-            else:
-                output_spec = tf.contrib.tpu.TPUEstimatorSpec(
-                    mode=mode,
-                    loss=total_loss,
-                    eval_metrics=eval_metrics,
-                    scaffold_fn=scaffold_fn)
+            # eval on single-gpu only
+            output_spec = tf.contrib.tpu.TPUEstimatorSpec(
+                mode=mode,
+                loss=total_loss,
+                eval_metrics=eval_metrics,
+                scaffold_fn=scaffold_fn)
         else:
-            if use_gpu and int(num_gpu_cores) >= 2:
-                output_spec = tf.estimator.EstimatorSpec(
-                    mode=mode,
-                    predictions={"probabilities": probabilities})
-            else:
-                output_spec = tf.contrib.tpu.TPUEstimatorSpec(
-                    mode=mode,
-                    predictions={"probabilities": probabilities},
-                    scaffold_fn=scaffold_fn)
+            # predict on single-gpu only
+            output_spec = tf.contrib.tpu.TPUEstimatorSpec(
+                mode=mode,
+                predictions={"probabilities": probabilities},
+                scaffold_fn=scaffold_fn)
 
         return output_spec
 
@@ -918,76 +922,70 @@ def main(_):
             FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
 
     is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
-    if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
-        tf.logging.info("Use normal RunConfig")
-        # https://github.com/tensorflow/tensorflow/issues/21470#issuecomment-422506263
-        dist_strategy = tf.contrib.distribute.MirroredStrategy(
-            num_gpus=FLAGS.num_gpu_cores,
-            cross_device_ops=AllReduceCrossDeviceOps('nccl', num_packs=FLAGS.num_gpu_cores),
-            # cross_device_ops=AllReduceCrossDeviceOps('hierarchical_copy'),
-        )
-        log_every_n_steps = 8
-        run_config = RunConfig(
-            train_distribute=dist_strategy,
-            eval_distribute=dist_strategy,
-            log_step_count_steps=log_every_n_steps,
-            model_dir=FLAGS.output_dir,
-            save_checkpoints_steps=FLAGS.save_checkpoints_steps)
-    else:
-        tf.logging.info("Use TPURunConfig")
-        run_config = tf.contrib.tpu.RunConfig(
-            cluster=tpu_cluster_resolver,
-            master=FLAGS.master,
-            model_dir=FLAGS.output_dir,
-            save_checkpoints_steps=FLAGS.save_checkpoints_steps,
-            tpu_config=tf.contrib.tpu.TPUConfig(
-                iterations_per_loop=FLAGS.iterations_per_loop,
-                num_shards=FLAGS.num_tpu_cores,
-                per_host_input_for_training=is_per_host))
-
-    train_examples = None
-    num_train_steps = None
-    num_warmup_steps = None
+
+    # https://github.com/tensorflow/tensorflow/issues/21470#issuecomment-422506263
+    dist_strategy = tf.contrib.distribute.MirroredStrategy(
+        num_gpus=FLAGS.num_gpu_cores,
+        cross_device_ops=AllReduceCrossDeviceOps('nccl', num_packs=FLAGS.num_gpu_cores),
+        # cross_device_ops=AllReduceCrossDeviceOps('hierarchical_copy'),
+    )
+    log_every_n_steps = 8
+    dist_run_config = RunConfig(
+        train_distribute=dist_strategy,
+        eval_distribute=dist_strategy,
+        log_step_count_steps=log_every_n_steps,
+        model_dir=FLAGS.output_dir,
+        save_checkpoints_steps=FLAGS.save_checkpoints_steps)
+
+    tpu_run_config = tf.contrib.tpu.RunConfig(
+        cluster=tpu_cluster_resolver,
+        master=FLAGS.master,
+        model_dir=FLAGS.output_dir,
+        save_checkpoints_steps=FLAGS.save_checkpoints_steps,
+        tpu_config=tf.contrib.tpu.TPUConfig(
+            iterations_per_loop=FLAGS.iterations_per_loop,
+            num_shards=FLAGS.num_tpu_cores,
+            per_host_input_for_training=is_per_host))
+
+    num_train_steps = 0
+    num_warmup_steps = 0
+    init_checkpoint = FLAGS.init_checkpoint
+    is_multi_gpu = FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2
     if FLAGS.do_train:
         train_examples = processor.get_train_examples(FLAGS.data_dir)
         num_train_steps = int(
             len(train_examples) / FLAGS.train_batch_size * FLAGS.num_train_epochs)
         num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
 
-    init_checkpoint = FLAGS.init_checkpoint
-
-    model_fn = model_fn_builder(
-        bert_config=bert_config,
-        num_labels=len(label_list),
-        init_checkpoint=init_checkpoint,
-        learning_rate=FLAGS.learning_rate,
-        num_train_steps=num_train_steps,
-        num_warmup_steps=num_warmup_steps,
-        use_tpu=FLAGS.use_tpu,
-        use_one_hot_embeddings=FLAGS.use_tpu,
-        use_gpu=FLAGS.use_gpu,
-        num_gpu_cores=FLAGS.num_gpu_cores,
-        fp16=FLAGS.use_fp16)
-
-    # If TPU is not available, this will fall back to normal Estimator on CPU
-    # or GPU.
-    if FLAGS.use_gpu and int(FLAGS.num_gpu_cores) >= 2:
-        tf.logging.info("Use normal Estimator")
-        estimator = Estimator(
-            model_fn=model_fn,
-            params={},
-            config=run_config)
-    else:
-        tf.logging.info("Use TPUEstimator")
-        estimator = tf.contrib.tpu.TPUEstimator(
+        model_fn = model_fn_builder(
+            bert_config=bert_config,
+            num_labels=len(label_list),
+            init_checkpoint=init_checkpoint,
+            learning_rate=FLAGS.learning_rate,
+            num_train_steps=num_train_steps,
+            num_warmup_steps=num_warmup_steps,
             use_tpu=FLAGS.use_tpu,
-            model_fn=model_fn,
-            config=run_config,
-            train_batch_size=FLAGS.train_batch_size,
-            eval_batch_size=FLAGS.eval_batch_size,
-            predict_batch_size=FLAGS.predict_batch_size)
+            use_one_hot_embeddings=FLAGS.use_tpu,
+            use_gpu=FLAGS.use_gpu,
+            num_gpu_cores=FLAGS.num_gpu_cores,
+            fp16=FLAGS.use_fp16)
+
+        # If TPU is not available, this will fall back to normal Estimator on CPU
+        # or GPU.
+        if is_multi_gpu:
+            estimator = Estimator(
+                model_fn=model_fn,
+                params={},
+                config=dist_run_config)
+        else:
+            estimator = tf.contrib.tpu.TPUEstimator(
+                use_tpu=FLAGS.use_tpu,
+                model_fn=model_fn,
+                config=tpu_run_config,
+                train_batch_size=FLAGS.train_batch_size,
+                eval_batch_size=FLAGS.eval_batch_size,
+                predict_batch_size=FLAGS.predict_batch_size)
 
-    if FLAGS.do_train:
         train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
         file_based_convert_examples_to_features(
             train_examples, label_list, FLAGS.max_seq_length, tokenizer, train_file)
@@ -1002,6 +1000,41 @@ def main(_):
             drop_remainder=True,
             batch_size=FLAGS.train_batch_size)
         estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
+        # TF Serving
+        if FLAGS.save_for_serving:
+            serving_dir = os.path.join(FLAGS.output_dir, 'serving')
+            save_for_serving(estimator, serving_dir, FLAGS.max_seq_length, not is_multi_gpu)
+
+        # Find the latest checkpoint
+        max_idx = 0
+        for filename in os.listdir(FLAGS.output_dir):
+            if filename.startswith('model.ckpt-'):
+                max_idx = max(int(filename.split('.')[1].split('-')[1]), max_idx)
+        init_checkpoint = os.path.join(FLAGS.output_dir, f'model.ckpt-{max_idx}')
+
+    if not FLAGS.do_eval and not FLAGS.do_predict:
+        return
+
+    model_fn = model_fn_builder(
+        bert_config=bert_config,
+        num_labels=len(label_list),
+        init_checkpoint=init_checkpoint,
+        learning_rate=FLAGS.learning_rate,
+        num_train_steps=num_train_steps,
+        num_warmup_steps=num_warmup_steps,
+        use_tpu=FLAGS.use_tpu,
+        use_one_hot_embeddings=FLAGS.use_tpu,
+        use_gpu=FLAGS.use_gpu,
+        num_gpu_cores=FLAGS.num_gpu_cores,
+        fp16=FLAGS.use_fp16)
+
+    estimator = tf.contrib.tpu.TPUEstimator(
+        use_tpu=FLAGS.use_tpu,
+        model_fn=model_fn,
+        config=tpu_run_config,
+        train_batch_size=FLAGS.train_batch_size,
+        eval_batch_size=FLAGS.eval_batch_size,
+        predict_batch_size=FLAGS.predict_batch_size)
 
     if FLAGS.do_eval:
         eval_examples = processor.get_dev_examples(FLAGS.data_dir)
@@ -1050,6 +1083,22 @@ def main(_):
                 tf.logging.info("  %s = %s", key, str(result[key]))
                 writer.write("%s = %s\n" % (key, str(result[key])))
 
+        # dump result as json file (easy parsing for other tasks)
+        class ExtEncoder(json.JSONEncoder):
+            def default(self, obj):
+                if isinstance(obj, np.integer):
+                    return int(obj)
+                if isinstance(obj, np.floating):
+                    return float(obj)
+                if isinstance(obj, np.ndarray):
+                    return obj.tolist()
+                else:
+                    return super(ExtEncoder, self).default(obj)
+
+        output_eval_file2 = os.path.join(FLAGS.output_dir, "eval_results.json")
+        with tf.gfile.GFile(output_eval_file2, "w") as writer:
+            json.dump(result, writer, indent=4, cls=ExtEncoder)
+
     if FLAGS.do_predict:
         predict_examples = processor.get_test_examples(FLAGS.data_dir)
         num_actual_predict_examples = len(predict_examples)
@@ -1097,11 +1146,6 @@ def main(_):
                 num_written_lines += 1
         assert num_written_lines == num_actual_predict_examples
 
-    if FLAGS.do_train and FLAGS.save_for_serving:
-        serving_dir = os.path.join(FLAGS.output_dir, 'serving')
-        is_tpu_estimator = not FLAGS.use_gpu or int(FLAGS.num_gpu_cores) < 2
-        save_for_serving(estimator, serving_dir, FLAGS.max_seq_length, is_tpu_estimator)
-
 
 if __name__ == "__main__":
     flags.mark_flag_as_required("data_dir")
diff --git a/run_custom_classifier.sh b/run_custom_classifier.sh
