Merge branch 'hepj-test' into 'main'

增加transformer-xl模型代码

See merge request dcutoolkit/deeplearing/dlexamples_new!23

TensorFlow/NLP/transformer-xl-master/.idea/.gitignore

+# Default ignored files
+/shelf/
+/workspace.xml

TensorFlow/NLP/transformer-xl-master/.idea/inspectionProfiles/profiles_settings.xml

+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>
\ No newline at end of file

TensorFlow/NLP/transformer-xl-master/.idea/misc.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8" project-jdk-type="Python SDK" />
+</project>
\ No newline at end of file

TensorFlow/NLP/transformer-xl-master/.idea/modules.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/transformer-xl-master.iml" filepath="$PROJECT_DIR$/.idea/transformer-xl-master.iml" />
+    </modules>
+  </component>
+</project>
\ No newline at end of file

TensorFlow/NLP/transformer-xl-master/.idea/transformer-xl-master.iml

+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>
\ No newline at end of file

TensorFlow/NLP/transformer-xl-master/README.md

+参考网址https://github.com/kimiyoung/transformer-xl
+
+# 1.数据准备
+
+数据集下载地址：http://mattmahoney.net/dc/enwik8.zip
+
+数据处理：在enwik8.zip同一文件夹下
+
+```
+python3 prep_enwik8.py
+```
+
+# 2.数据转换
+
+1.修改enwik8_base_gpu.sh脚本中DATA_ROOT为数据路径
+
+2.执行
+
+```
+bash ./scripts/enwik8_base_gpu.sh train_data
+```
+
+# 3.模型运行
+
+```
+#单卡运行
+bash ./scripts/enwik8_base1_gpu.sh train
+#多卡运行
+bash scripts/enwik8_base_gpu.sh
+```
+
+# 4.重要参数说明
+
+```
+enwik8_base_gpu.sh & enwik8_base1_gpu.sh
+--MODEL_DIR   --模型保存位置
+--NUM_CORE    --为使用GPU数量，
+--BSZ         --batchs_size（注意使用数据转化是使用的BSZ和训练使用的BSZ需要保持一致）
+--TRAIN_STEPS --训练步数
+```
+

TensorFlow/NLP/transformer-xl-master/README.old.md

+
+## Introduction
+
+This directory contains our TF implementation of Transformer-XL. Note that our state-of-the-art results reported in the paper were obtained by training the model on a large-scale TPU cluster, and our gpu codebase currently does not support distributed training. Here we provide two sets of hyperparameters and scripts:
+- `*large_tpu.sh` are for the SoTA setting on TPUs. These are exactly the commands we used to obtained our best results.
+- `*base_gpu.sh` are for the base models which can be run on a few GPUs.
+
+
+## Prerequisite
+
+- Python 2.7
+- Tensorflow [1.12.0](https://github.com/tensorflow/tensorflow/releases/tag/v1.12.0)
+
+
+
+## Obtain and evaluate pretrained SoTA models
+
+#### 1. Download preprocessed data (vocab) & pretrained models
+
+(a) Set your own `DATA_ROOT` in `sota/download.sh` (default to `./`), which will be the root diretory of downloaded model.
+
+(b) Then, download the model & data by `bash sota/download.sh`. After downloading, the expected directory structure is as follows
+
+```markdown
+pretrained_xl
+  tf_enwik8/
+    data/
+      cache.pkl
+      corpus-info.json
+    model/
+      checkpoint
+      model.ckpt*
+  tf_wt103/
+  	...
+  ...
+```
+
+**Note**: we include preprocessed data in the download files to make sure the **same vocabulary** is used. Please see the code `tf/data_utils.py` to understand the data structure.
+
+
+
+#### 2. Run evaluation scripts to replicate SoTA results on GPUs
+
+- **enwik8**: modify the script `sota/enwik8.sh` accordingly (see below)
+  - set `DATA_ROOT` to the same folder used in the download step (default to `./`)
+  - set `TEST_NUM_CORE ` (number of GPUs to use): we recommend 2 GPUs => about 60 mins
+  - run the script: `bash sota/enwik8.sh`
+
+- **lm1b**: modify the script `sota/lm1b.sh` accordingly  (see below)
+  - set `DATA_ROOT` to the same folder used in the download step (default to `./`)
+  - set `TEST_NUM_CORE ` (number of GPUs to use): we recommend 1 GPUs => less than 5 mins
+  - run the script: `bash sota/lm1b.sh`
+
+- **wt103**:  modify the script `sota/wt103.sh` accordingly  (see below)
+  - set `DATA_ROOT` to the same folder used in the download step (default to `./`)
+  - set `TEST_NUM_CORE ` (number of GPUs to use): we recommend 1 GPUs => less than 5 mins
+  - run the script: `bash sota/wt103.sh`
+
+- **text8**:  modify the script `sota/text8.sh` accordingly  (see below)
+  - set `DATA_ROOT` to the same folder used in the download step (default to `./`)
+  - set `TEST_NUM_CORE ` (number of GPUs to use): we recommend 2 GPUs => about 60 mins
+  - run the script: `bash sota/text8.sh`
+
+
+#### 3. Resources Needed for SoTA Model Training
+
+We used 32, 32, 64, and 512 TPU cores for training our best models on enwik8, text8, wt103, and lm1b respectively. The training time for each model ranges from 2 to 5 days.
+
+
+
+## Train "Transformer-XL" from scratch with GPUs or TPUs
+
+### 1. Download raw data
+
+`bash getdata.sh`
+
+
+
+### 2. Preprocess, training and evaluation
+
+For `dataset` in `[enwik8, lm1b, wt103, text8]`:
+
+- check out `scripts/dataset_base_gpu.sh` for GPU training and evaluation
+- check out `scripts/dataset_large_tpu.sh` for TPU training and evaluation
+
+
+
+#### (1) Preprocess raw data and create tfrecords
+
+**NOTE**: The preprocessing for GPU and TPU are different. So, you have to run them separately.
+
+GPU:
+
+- create training and validation data: `bash scripts/dataset_bas_gpu.sh train_data`
+- create test data: `bash scripts/dataset_base_gpu.sh test_data`
+
+TPU:
+
+- Set the Google storage URL  in `scripts/dataset_large_tpu.sh`:
+  - `GSDATA`: data URL
+  - `GSEXP`: experiment URL
+- create training and validation data: `bash scripts/dataset_large_tpu.sh train_data`
+- create test data: `bash scripts/dataset_large_tpu.sh test_data`
+
+
+
+#### (2) Run training
+
+Base models on GPUs:
+
+- Modify the configurations in `scripts/dataset_base_gpu.sh`  according to your needs.
+- `bash scripts/dataset_base_gpu.sh train`
+- If enough resources are available, increasing the model sizes (e.g., `N_LAYER`, `D_MODEL`, `D_EMBED`, `D_HEAD`, `D_INNER`) so that they are closer to the values defined in `scripts/dataset_large_tpu.sh`. Likewise, when resources are limited, decrease the model sizes. It is recommended to ensure that `D_MODEL == D_EMBED` and `D_MODEL == N_HEAD x D_HEAD`. When the model sizes increase, remember to increase `warmup_steps` accordingly to alleviate optimization difficulties.
+- Adjust the `NUM_CORE` parameter to reflect the number of GPUs to use.
+
+Larger models on TPUs:
+
+- Modify the configurations in `scripts/dataset_large_tpu.sh`  according to your needs.
+- `bash scripts/dataset_large_tpu.sh train`
+
+
+
+#### (3) Run evaluation
+
+Base models on GPUs:
+
+- `bash scripts/dataset_base_gpu.sh eval --eval_ckpt_path PATH_TO_CKPT`
+
+Larger models on TPUs:
+
+- `bash scripts/dataset_base_tpu.sh eval --eval_ckpt_path PATH_TO_CKPT`

TensorFlow/NLP/transformer-xl-master/avg_checkpoints.py

+# coding=utf-8
+# Copyright 2018 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Script to average values of variables in a list of checkpoint files."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import numpy as np
+import six
+from six.moves import zip  # pylint: disable=redefined-builtin
+import tensorflow as tf
+
+flags = tf.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("checkpoints", "",
+                    "Comma-separated list of checkpoints to average.")
+flags.DEFINE_integer("num_last_checkpoints", 0,
+                     "Averages the last N saved checkpoints."
+                     " If the checkpoints flag is set, this is ignored.")
+flags.DEFINE_string("prefix", "",
+                    "Prefix (e.g., directory) to append to each checkpoint.")
+flags.DEFINE_string("output_path", "/tmp/averaged.ckpt",
+                    "Path to output the averaged checkpoint to.")
+
+
+def checkpoint_exists(path):
+  return (tf.gfile.Exists(path) or tf.gfile.Exists(path + ".meta") or
+          tf.gfile.Exists(path + ".index"))
+
+
+def main(_):
+  tf.logging.set_verbosity(tf.logging.INFO)
+  if FLAGS.checkpoints:
+    # Get the checkpoints list from flags and run some basic checks.
+    checkpoints = [c.strip() for c in FLAGS.checkpoints.split(",")]
+    checkpoints = [c for c in checkpoints if c]
+    if not checkpoints:
+      raise ValueError("No checkpoints provided for averaging.")
+    if FLAGS.prefix:
+      checkpoints = [FLAGS.prefix + c for c in checkpoints]
+  else:
+    assert FLAGS.num_last_checkpoints >= 1, "Must average at least one model"
+    assert FLAGS.prefix, ("Prefix must be provided when averaging last"
+                          " N checkpoints")
+    checkpoint_state = tf.train.get_checkpoint_state(
+        os.path.dirname(FLAGS.prefix))
+    # Checkpoints are ordered from oldest to newest.
+    checkpoints = checkpoint_state.all_model_checkpoint_paths[
+        -FLAGS.num_last_checkpoints:]
+
+  checkpoints = [c for c in checkpoints if checkpoint_exists(c)]
+  if not checkpoints:
+    if FLAGS.checkpoints:
+      raise ValueError(
+          "None of the provided checkpoints exist. %s" % FLAGS.checkpoints)
+    else:
+      raise ValueError("Could not find checkpoints at %s" %
+                       os.path.dirname(FLAGS.prefix))
+
+  # Read variables from all checkpoints and average them.
+  tf.logging.info("Reading variables and averaging checkpoints:")
+  for c in checkpoints:
+    tf.logging.info("%s ", c)
+  var_list = tf.contrib.framework.list_variables(checkpoints[0])
+  var_values, var_dtypes = {}, {}
+  for (name, shape) in var_list:
+    if not name.startswith("global_step"):
+      var_values[name] = np.zeros(shape)
+  for checkpoint in checkpoints:
+    reader = tf.contrib.framework.load_checkpoint(checkpoint)
+    for name in var_values:
+      tensor = reader.get_tensor(name)
+      var_dtypes[name] = tensor.dtype
+      var_values[name] += tensor
+    tf.logging.info("Read from checkpoint %s", checkpoint)
+  for name in var_values:  # Average.
+    var_values[name] /= len(checkpoints)
+
+  with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
+    tf_vars = [
+        tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
+        for v in var_values
+    ]
+  placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
+  assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
+  global_step = tf.Variable(
+      0, name="global_step", trainable=False, dtype=tf.int64)
+  saver = tf.train.Saver(tf.all_variables())
+
+  # Build a model consisting only of variables, set them to the average values.
+  with tf.Session() as sess:
+    sess.run(tf.initialize_all_variables())
+    for p, assign_op, (name, value) in zip(placeholders, assign_ops,
+                                           six.iteritems(var_values)):
+      sess.run(assign_op, {p: value})
+    # Use the built saver to save the averaged checkpoint.
+    saver.save(sess, FLAGS.output_path, global_step=global_step)
+
+  tf.logging.info("Averaged checkpoints saved in %s", FLAGS.output_path)
+
+
+if __name__ == "__main__":
+  tf.app.run()

TensorFlow/NLP/transformer-xl-master/data_utils.py

+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import os
+from functools import partial
+
+from collections import Counter, OrderedDict
+import pickle
+import json
+import multiprocessing as mp
+
+import numpy as np
+
+from absl import flags
+import tensorflow as tf
+from vocabulary import Vocab
+
+from tensorflow.gfile import Exists as exists
+from tensorflow.gfile import MakeDirs as makedirs
+from tensorflow.gfile import Glob as glob
+
+
+def _preprocess(shard, train, vocab, save_dir, cutoffs, bin_sizes, bsz, tgt_len,
+                num_core_per_host, use_tpu, num_shuffle):
+  file_names = []
+  num_batch = 0
+
+  path = train[shard]
+  data_shard = vocab.encode_file(path, ordered=False, add_double_eos=True)
+
+  for shuffle in range(num_shuffle):
+    basename = "train-{:03d}-{:02d}".format(shard, shuffle)
+    print("Processing shard {} shuffle {}".format(shard, shuffle))
+
+    np.random.shuffle(data_shard)
+    file_name, num_batch_shuffle = create_ordered_tfrecords(
+        save_dir, basename, np.concatenate(data_shard), bsz, tgt_len,
+        num_core_per_host, cutoffs, bin_sizes, use_tpu=use_tpu)
+    file_names.append(file_name)
+    num_batch += num_batch_shuffle
+
+  return file_names, num_batch
+
+
+class Corpus(object):
+  def __init__(self, path, dataset, *args, **kwargs):
+    self.dataset = dataset
+    self.vocab = Vocab(*args, **kwargs)
+
+    if self.dataset in ["ptb", "wt2", "enwik8", "text8"]:
+      self.vocab.count_file(os.path.join(path, "train.txt"))
+      self.vocab.count_file(os.path.join(path, "valid.txt"))
+      self.vocab.count_file(os.path.join(path, "test.txt"))
+    elif self.dataset == "wt103":
+      self.vocab.count_file(os.path.join(path, "train.txt"))
+    elif self.dataset == "lm1b":
+      train_path_pattern = os.path.join(
+          path, "1-billion-word-language-modeling-benchmark-r13output",
+          "training-monolingual.tokenized.shuffled", "news.en-*")
+      train_paths = glob(train_path_pattern)
+
+      # the vocab will load from file when build_vocab() is called
+      # for train_path in sorted(train_paths):
+      #   self.vocab.count_file(train_path, verbose=True)
+
+    self.vocab.build_vocab()
+
+    if self.dataset in ["ptb", "wt2", "wt103"]:
+      self.train = self.vocab.encode_file(
+          os.path.join(path, "train.txt"), ordered=True)
+      self.valid = self.vocab.encode_file(
+          os.path.join(path, "valid.txt"), ordered=True)
+      self.test  = self.vocab.encode_file(
+          os.path.join(path, "test.txt"), ordered=True)
+    elif self.dataset in ["enwik8", "text8"]:
+      self.train = self.vocab.encode_file(
+          os.path.join(path, "train.txt"), ordered=True, add_eos=False)
+      self.valid = self.vocab.encode_file(
+          os.path.join(path, "valid.txt"), ordered=True, add_eos=False)
+      self.test  = self.vocab.encode_file(
+          os.path.join(path, "test.txt"), ordered=True, add_eos=False)
+    elif self.dataset == "lm1b":
+      self.train = train_paths
+      valid_path = os.path.join(path, "valid.txt")
+      test_path = valid_path
+      self.valid = self.vocab.encode_file(
+          valid_path, ordered=True, add_double_eos=True)
+      self.test  = self.vocab.encode_file(
+          test_path, ordered=True, add_double_eos=True)
+
+    if self.dataset == "wt103":
+      self.cutoffs = [0, 20000, 40000, 200000] + [len(self.vocab)]
+    elif self.dataset == "lm1b":
+      self.cutoffs = [0, 60000, 100000, 640000] + [len(self.vocab)]
+    else:
+      self.cutoffs = []
+
+
+  def convert_to_tfrecords(self, split, save_dir, bsz, tgt_len,
+                           num_core_per_host, **kwargs):
+    FLAGS = kwargs.get('FLAGS')
+
+    file_names = []
+    use_tpu = FLAGS.use_tpu and not (split == "test" and num_core_per_host == 1)
+
+    if use_tpu:
+      record_name = "record_info-{}.bsz-{}.tlen-{}.core-{}.json".format(
+           split, bsz, tgt_len, num_core_per_host)
+    else:
+      record_name = "record_info-{}.bsz-{}.tlen-{}.json".format(
+           split, bsz, tgt_len)
+
+    record_info_path = os.path.join(save_dir, record_name)
+
+    if self.dataset in ["ptb", "wt2", "wt103", "enwik8", "text8"]:
+      data = getattr(self, split)
+      bin_sizes = get_bin_sizes(
+          data, bsz // num_core_per_host, tgt_len, self.cutoffs)
+      file_name, num_batch = create_ordered_tfrecords(
+          save_dir, split, data, bsz, tgt_len, num_core_per_host,
+          self.cutoffs, bin_sizes,
+          num_passes=FLAGS.num_passes if split == 'train' and use_tpu else 1,
+          use_tpu=use_tpu)
+      file_names.append(file_name)
+    elif self.dataset == "lm1b":
+      bin_sizes = get_bin_sizes(
+          self.valid, bsz // num_core_per_host, tgt_len, self.cutoffs)
+      if split == "train":
+        np.random.seed(123456)
+        num_batch = 0
+
+        if FLAGS.num_procs > 1:
+          _preprocess_wrapper = partial(_preprocess,
+              train=self.train, vocab=self.vocab, save_dir=save_dir,
+              cutoffs=self.cutoffs, bin_sizes=bin_sizes, bsz=bsz,
+              tgt_len=tgt_len, num_core_per_host=num_core_per_host,
+              use_tpu=use_tpu, num_shuffle=FLAGS.num_shuffle)
+
+          pool = mp.Pool(processes=FLAGS.num_procs)
+          results = pool.map(_preprocess_wrapper, range(len(self.train)))
+          for res in results:
+            file_names.extend(res[0])
+            num_batch += res[1]
+        else:
+          for shard, path in enumerate(self.train):
+            data_shard = self.vocab.encode_file(path, ordered=False,
+                                                add_double_eos=True)
+
+            num_shuffle = FLAGS.num_shuffle
+
+            for shuffle in range(num_shuffle):
+              print("Processing shard {} shuffle {}".format(shard, shuffle))
+              basename = "train-{:03d}-{:02d}".format(shard, shuffle)
+              np.random.shuffle(data_shard)
+              file_name, num_batch_ = create_ordered_tfrecords(
+                  save_dir, basename, np.concatenate(data_shard), bsz, tgt_len,
+                  num_core_per_host,
+                  self.cutoffs, bin_sizes, use_tpu=use_tpu)
+              file_names.append(file_name)
+              num_batch += num_batch_
+
+      else:
+        file_name, num_batch = create_ordered_tfrecords(
+            save_dir, split, getattr(self, split), bsz, tgt_len,
+            num_core_per_host,
+            self.cutoffs, bin_sizes, use_tpu=use_tpu)
+        file_names.append(file_name)
+
+    with open(record_info_path, "w") as fp:
+      record_info = {
+        "filenames": file_names,
+        "bin_sizes": bin_sizes,
+        "num_batch": num_batch
+      }
+      json.dump(record_info, fp)
+
+
+def get_bin_sizes(data, batch_size, tgt_len, cutoffs, std_mult=[2.5, 2.5, 2.5]):
+  """
+    Note: the `batch_size` here should be per-core batch size
+  """
+  bin_sizes = []
+
+  def _nearest_to_eight(x): # so that it's faster on TPUs
+    y = x - x % 8
+    return y + 8 if x % 8 >= 4 else max(8, y)
+
+  if cutoffs:
+    num_batch = len(data) // batch_size // tgt_len
+
+    data = data[:batch_size * num_batch * tgt_len]
+    data = data.reshape(batch_size, num_batch, tgt_len)
+
+    tot = batch_size * tgt_len
+    for b, (left, right) in enumerate(zip(cutoffs[1:-1], cutoffs[2:])):
+      mask = (data >= left) * (data < right)
+      percents = mask.astype(np.float64).sum(2).sum(0) / tot
+      mean = np.mean(percents)
+      std = np.std(percents)
+
+      bin_size = int(math.ceil(tgt_len * batch_size * (mean + std_mult[b] * std)))
+      bin_size = _nearest_to_eight(bin_size)
+      bin_sizes.append(bin_size)
+
+  return bin_sizes
+
+
+def _int64_feature(values):
+  return tf.train.Feature(int64_list=tf.train.Int64List(value=values))
+
+def _float_feature(values):
+  return tf.train.Feature(float_list=tf.train.FloatList(value=values))
+
+def batchify(data, batch_size, num_passes):
+  """
+    if use_tpu = True: num_passes > 1 
+    
+    Since TPU training requires entire [bsz x tgt_len] chunks, it can discard
+    as many as `bsz * tgt_len` tokens in training. When `bsz` and `tgt_len` are 
+    both large, as in the case of TPU training for Transformer-XL, the problem
+    may lead to detectable performance drop. 
+
+    Here, we use multiple randomly shifted copies to deal with this problem.
+  """
+  if num_passes > 1:
+    data_len = len(data)
+    double_data = np.concatenate([data, data])
+    data_list = []
+    for i in range(num_passes):
+      start = np.random.randint(0, data_len)
+      data_list.append(double_data[start:start+data_len])
+    data = np.concatenate(data_list)
+
+  num_step = len(data) // batch_size
+  data = data[:batch_size * num_step]
+  data = data.reshape(batch_size, num_step)
+
+  return data
+
+
+def create_ordered_tfrecords(save_dir, basename, data, batch_size, tgt_len,
+                             num_core_per_host, cutoffs=[], bin_sizes=[], 
+                             num_passes=1, use_tpu=False):
+
+  if use_tpu:
+    file_name = "{}.bsz-{}.tlen-{}.core-{}.tfrecords".format(
+        basename, batch_size, tgt_len, num_core_per_host)
+  else:
+    file_name = "{}.bsz-{}.tlen-{}.tfrecords".format(
+        basename, batch_size, tgt_len)
+
+  save_path = os.path.join(save_dir, file_name)
+  record_writer = tf.python_io.TFRecordWriter(save_path)
+
+  batched_data = batchify(data, batch_size, num_passes)
+
+  num_batch = 0
+  # for t in range(0, batched_data.shape[1] - tgt_len - 1, tgt_len):
+  for t in range(0, batched_data.shape[1] - 1, tgt_len):
+    cur_tgt_len = min(batched_data.shape[1] - 1 - t, tgt_len)
+    # drop the remainder if use tpu
+    if use_tpu and cur_tgt_len < tgt_len: 
+      break
+    if num_batch % 500 == 0:
+      print("  processing batch {}".format(num_batch))
+    for idx in range(batch_size):
+      inputs = batched_data[idx, t:t + cur_tgt_len]
+      labels = batched_data[idx, t + 1:t + cur_tgt_len + 1]
+
+      # features dict
+      feature = {
+          "inputs": _int64_feature(inputs),
+          "labels": _int64_feature(labels),
+      }
+
+      if len(cutoffs) > 0 and use_tpu:
+        # validate `bin_sizes` and `cutoffs`
+        assert len(cutoffs) - len(bin_sizes) == 2, \
+          "len(cutoffs) - len(bin_sizes) != 2"
+
+        # mask for bin 0
+        left, right = cutoffs[:2]
+        inp_mask = ((inputs >= left) * (inputs < right)).astype(np.float32)
+        tgt_mask = ((labels >= left) * (labels < right)).astype(np.float32)
+
+        feature["inp_mask"] = _float_feature(inp_mask)
+        feature["tgt_mask"] = _float_feature(tgt_mask)
+
+        # refresh `inp_cnts` and `tgt_cnts` for each TPU core
+        if idx % (batch_size // num_core_per_host) == 0:
+          inp_cnts = [0] * len(bin_sizes)
+          tgt_cnts = [0] * len(bin_sizes)
+
+        head_labels = np.copy(labels)
+        inp_pos_per_bin, tgt_pos_per_bin = [], []
+        for b, (left, right) in enumerate(zip(cutoffs[1:-1], cutoffs[2:])):
+          inp_pos = np.where((inputs >= left) * (inputs < right))[0]
+          tgt_pos = np.where((labels >= left) * (labels < right))[0]
+          inp_pos_per_bin.append(inp_pos)
+          tgt_pos_per_bin.append(tgt_pos)
+
+          head_labels[tgt_pos] = cutoffs[1] + b
+
+        feature["head_labels"] = _int64_feature(head_labels)
+
+        # permutation feature
+        def _add_perm_feature(feature, pos_per_bin, cnts, prefix):
+          for b, pos in enumerate(pos_per_bin):
+            idx_tuple = []
+            for p in pos:
+              if cnts[b] < bin_sizes[b]:
+                idx_tuple.append([p, cnts[b]])
+                cnts[b] += 1
+              else:
+                break
+
+            n_tup = len(idx_tuple)
+            tup = np.array(idx_tuple).reshape(n_tup * 2)
+
+            feature["{}_cnt_{}".format(prefix, b)] = _int64_feature([n_tup])
+            feature["{}_tup_{}".format(prefix, b)] = _int64_feature(tup)
+
+        _add_perm_feature(feature, inp_pos_per_bin, inp_cnts, "inp")
+        _add_perm_feature(feature, tgt_pos_per_bin, tgt_cnts, "tgt")
+
+      example = tf.train.Example(features=tf.train.Features(feature=feature))
+      record_writer.write(example.SerializeToString())
+
+    num_batch += 1
+
+  record_writer.close()
+  print("Done writing {}. batches: {}".format(file_name, num_batch))
+
+  return file_name, num_batch
+
+
+def get_lm_corpus(data_dir, dataset):
+  fn = os.path.join(data_dir, "cache.pkl")
+
+  if exists(fn):
+    print("Loading cached dataset...")
+    with open(fn, "rb") as fp:
+      corpus = pickle.load(fp)
+  else:
+    print("Producing dataset...")
+    kwargs = {}
+    if dataset in ["wt103", "wt2"]:
+      kwargs["special"] = ["<eos>"]
+      kwargs["lower_case"] = False
+    elif dataset == "ptb":
+      kwargs["special"] = ["<eos>"]
+      kwargs["lower_case"] = True
+    elif dataset == "lm1b":
+      kwargs["special"] = []
+      kwargs["lower_case"] = False
+      kwargs["vocab_file"] = os.path.join(data_dir, "1b_word_vocab.txt")
+    elif dataset in ["enwik8", "text8"]:
+      pass
+
+    corpus = Corpus(data_dir, dataset, **kwargs)
+
+    print("Saving dataset...")
+    with open(fn, "wb") as fp:
+      pickle.dump(corpus, fp, protocol=2)
+
+    corpus_info = {
+      "vocab_size" : len(corpus.vocab),
+      "cutoffs" : corpus.cutoffs,
+      "dataset" : corpus.dataset
+    }
+    with open(os.path.join(data_dir, "corpus-info.json"), "w") as fp:
+      json.dump(corpus_info, fp)
+
+  return corpus
+
+
+def main(unused_argv):
+  del unused_argv  # Unused
+
+  corpus = get_lm_corpus(FLAGS.data_dir, FLAGS.dataset)
+
+  save_dir = os.path.join(FLAGS.data_dir, "tfrecords")
+  if not exists(save_dir):
+    makedirs(save_dir)
+
+  # test mode
+  if FLAGS.per_host_test_bsz > 0:
+    corpus.convert_to_tfrecords("test", save_dir, FLAGS.per_host_test_bsz,
+                                FLAGS.tgt_len, FLAGS.num_core_per_host, 
+                                FLAGS=FLAGS)
+    return
+
+  for split, batch_size in zip(
+      ["train", "valid"],
+      [FLAGS.per_host_train_bsz, FLAGS.per_host_valid_bsz]):
+
+    if batch_size <= 0: continue
+    print("Converting {} set...".format(split))
+    corpus.convert_to_tfrecords(split, save_dir, batch_size, FLAGS.tgt_len,
+                                FLAGS.num_core_per_host, FLAGS=FLAGS)
+
+
+def load_record_info(record_info_dir, split, per_host_bsz, tgt_len,
+                     num_core_per_host, use_tpu):
+  if use_tpu:
+    record_name = "record_info-{}.bsz-{}.tlen-{}.core-{}.json".format(
+        split, per_host_bsz, tgt_len, num_core_per_host)
+  else:
+    record_name = "record_info-{}.bsz-{}.tlen-{}.json".format(
+        split, per_host_bsz, tgt_len)
+
+  record_info_path = os.path.join(record_info_dir, record_name)
+  #改
+  print("="*50)
+  print(record_info_path)
+  print("="*50)
+  with open(record_info_path, "r") as fp:
+    record_info = json.load(fp)
+
+  return record_info
+
+def get_input_fn(record_info_dir, split, per_host_bsz, tgt_len,
+                 num_core_per_host, num_hosts=1, use_tpu=False):
+  """Creates input function."""
+  record_info = load_record_info(record_info_dir, split, per_host_bsz, tgt_len,
+                                 num_core_per_host, use_tpu=use_tpu)
+
+  file_names = record_info["filenames"]
+  bin_sizes = record_info["bin_sizes"]
+  num_batch = record_info["num_batch"]
+
+  tf.logging.info("[{}] File names {}".format(split, file_names))
+
+  def input_fn(params):
+    # per-core batch size
+    per_core_bsz = params["batch_size"]
+
+    # data_dir could be a remote path, e.g., a google storage url
+    data_dir = params["data_dir"]
+
+    def parser(record):
+      # preprocess "inp_perm" and "tgt_perm"
+      def _process_perm_feature(example, prefix):
+        for b in range(len(bin_sizes)):
+          cnt = example.pop("{}_cnt_{}".format(prefix, b))[0]
+          tup = example.pop("{}_tup_{}".format(prefix, b))
+
+          tup = tf.reshape(
+              tf.sparse_tensor_to_dense(tup),
+              shape=[cnt, 2])
+
+          # tf.float32
+          perm = tf.sparse_to_dense(
+              sparse_indices=tup,
+              output_shape=[tgt_len, bin_sizes[b]],
+              sparse_values=1.0,
+              default_value=0.0)
+
+          example["{}_perm_{}".format(prefix, b)] = perm
+
+      # whether allow the last batch with a potentially shorter length
+      if use_tpu:
+        record_spec = {
+            "inputs": tf.FixedLenFeature([tgt_len], tf.int64),
+            "labels": tf.FixedLenFeature([tgt_len], tf.int64),
+        }
+      else:
+        record_spec = {
+            "inputs": tf.VarLenFeature(tf.int64),
+            "labels": tf.VarLenFeature(tf.int64),
+        }
+
+      # permutation related features
+      if bin_sizes and use_tpu:
+        # tf.float32
+        record_spec["inp_mask"] = tf.FixedLenFeature([tgt_len], tf.float32)
+        record_spec["tgt_mask"] = tf.FixedLenFeature([tgt_len], tf.float32)
+
+        record_spec["head_labels"] = tf.FixedLenFeature([tgt_len], tf.int64)
+
+        for b in range(len(bin_sizes)):
+          record_spec["inp_cnt_{}".format(b)] = tf.FixedLenFeature([1], tf.int64)
+          record_spec["inp_tup_{}".format(b)] = tf.VarLenFeature(tf.int64)
+          record_spec["tgt_cnt_{}".format(b)] = tf.FixedLenFeature([1], tf.int64)
+          record_spec["tgt_tup_{}".format(b)] = tf.VarLenFeature(tf.int64)
+
+      # retrieve serialized example
+      example = tf.parse_single_example(
+          serialized=record,
+          features=record_spec)
+
+      # transform permutation tuples to permutation matrices
+      if bin_sizes and use_tpu:
+        _process_perm_feature(example, "inp")
+        _process_perm_feature(example, "tgt")
+
+      # cast int64 into int32
+      # cast sparse to dense
+      for key in list(example.keys()):
+        val = example[key]
+        if tf.keras.backend.is_sparse(val):
+          val = tf.sparse.to_dense(val)
+        if val.dtype == tf.int64:
+          val = tf.to_int32(val)
+        example[key] = val
+
+      if use_tpu:
+        return example
+      else:
+        return example["inputs"], example["labels"]
+
+    file_paths = []
+    for file_name in file_names:
+      file_path = os.path.join(data_dir, file_name)
+      file_paths.append(file_path)
+
+    if split == "train":
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+      if len(file_paths) > 1:
+        dataset = dataset.shuffle(len(file_paths)).repeat()
+        dataset = tf.data.TFRecordDataset(dataset)
+      elif num_hosts > 1:
+        host_id = params["context"].current_host
+        # drop the remaining batches
+        num_batch_per_host = num_batch // num_hosts
+
+        my_start_sample_id = (host_id * num_batch_per_host * num_core_per_host *
+                              per_core_bsz)
+        my_sample_num = num_batch_per_host * num_core_per_host * per_core_bsz
+        dataset = tf.data.TFRecordDataset(dataset).skip(
+            my_start_sample_id).take(my_sample_num)
+      else:
+        dataset = tf.data.TFRecordDataset(dataset)
+
+      dataset = dataset.map(parser).cache().repeat()
+      dataset = dataset.batch(per_core_bsz, drop_remainder=True)
+      dataset = dataset.prefetch(num_core_per_host * per_core_bsz)
+    else:
+      # do not shuffle, repeat or cache in evaluation
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+      dataset = tf.data.TFRecordDataset(dataset)
+      dataset = dataset.map(parser)
+      dataset = dataset.batch(per_core_bsz, drop_remainder=True)
+
+    return dataset
+
+  if split == "train" and num_hosts > 1:
+    record_info["num_batch"] = num_batch // num_hosts
+
+  return input_fn, record_info
+
+def get_corpus_info(corpus_info_path):
+  with open(corpus_info_path, "r") as fp:
+    corpus_info = json.load(fp)
+  return corpus_info
+
+if __name__ == "__main__":
+  FLAGS = flags.FLAGS
+  flags.DEFINE_string("data_dir", None,
+        help="Location of the data corpus")
+  flags.DEFINE_enum("dataset", "wt103",
+        ["ptb", "wt2", "wt103", "lm1b", "enwik8", "text8"],
+        help="Dataset name.")
+  flags.DEFINE_integer("per_host_train_bsz", 60,
+        help="train batch size each host")
+  flags.DEFINE_integer("per_host_valid_bsz", 60,
+        help="valid batch size each host")
+  flags.DEFINE_integer("per_host_test_bsz", 0,
+        help="If > 0, enter test mode and process test set only."
+             "Otherwise, process train and dev sets only.")
+  flags.DEFINE_integer("tgt_len", 70,
+        help="number of tokens to predict")
+  flags.DEFINE_integer("max_batch", -1,
+        help="run in debug mode")
+  flags.DEFINE_integer("num_core_per_host", 8,
+        help="8 for TPU v2.")
+  flags.DEFINE_bool("debug", default=False,
+        help="Process only the first batch without shuffle for lm1b.")
+  flags.DEFINE_integer("num_procs", 1,
+        help="number of processes")
+  flags.DEFINE_integer("num_passes", 10,
+        help="number of passes when use_tpu=True")
+  flags.DEFINE_integer("num_shuffle", 4,
+        help="number of shuffles for lm1b")
+  flags.DEFINE_bool("use_tpu", True,
+        help="use tpu")
+
+  tf.app.run(main)

TensorFlow/NLP/transformer-xl-master/gpu_utils.py

+import os
+import tensorflow as tf
+
+def assign_to_gpu(gpu=0, ps_dev="/device:CPU:0"):
+    def _assign(op):
+        node_def = op if isinstance(op, tf.NodeDef) else op.node_def
+        if node_def.op == "Variable":
+            return ps_dev
+        else:
+            return "/gpu:%d" % gpu
+    return _assign
+
+
+def average_grads_and_vars(tower_grads_and_vars):
+    def average_dense(grad_and_vars):
+        if len(grad_and_vars) == 1:
+            return grad_and_vars[0][0]
+
+        grad = grad_and_vars[0][0]
+        for g, _ in grad_and_vars[1:]:
+            grad += g
+        return grad / len(grad_and_vars)
+
+    def average_sparse(grad_and_vars):
+        if len(grad_and_vars) == 1:
+            return grad_and_vars[0][0]
+
+        indices = []
+        values = []
+        for g, _ in grad_and_vars:
+            indices += [g.indices]
+            values += [g.values]
+        indices = tf.concat(indices, 0)
+        values = tf.concat(values, 0) / len(grad_and_vars)
+        return tf.IndexedSlices(values, indices, grad_and_vars[0][0].dense_shape)
+
+    average_grads_and_vars = []
+    for grad_and_vars in zip(*tower_grads_and_vars):
+        if grad_and_vars[0][0] is None:
+            grad = None
+        elif isinstance(grad_and_vars[0][0], tf.IndexedSlices):
+            grad = average_sparse(grad_and_vars)
+        else:
+            grad = average_dense(grad_and_vars)
+        # Keep in mind that the Variables are redundant because they are shared
+        # across towers. So .. we will just return the first tower's pointer to
+        # the Variable.
+        v = grad_and_vars[0][1]
+        grad_and_var = (grad, v)
+        average_grads_and_vars.append(grad_and_var)
+    return average_grads_and_vars
+
+
+def load_from_checkpoint(saver, logdir):
+    sess = tf.get_default_session()
+    ckpt = tf.train.get_checkpoint_state(logdir)
+    if ckpt and ckpt.model_checkpoint_path:
+        if os.path.isabs(ckpt.model_checkpoint_path):
+            # Restores from checkpoint with absolute path.
+            saver.restore(sess, ckpt.model_checkpoint_path)
+        else:
+            # Restores from checkpoint with relative path.
+            saver.restore(sess, os.path.join(logdir, ckpt.model_checkpoint_path))
+        return True
+    return False