added_training_stuff

2026-09-19 12:10:28 +08:00 · 2019-09-27 14:39:18 -06:00
parent ef45ad99d0
commit d71a2e95eb
14 changed files with 447998 additions and 1 deletions
@@ -0,0 +1,85 @@
+# Fine-Tuning the Model on Custom Dataset
+
+
+This folder contains sample code to fine-tune the model on custom data. It is primarily targeted for GPU usage, but there are pointers throughout showing how to run on TPUs as well. 
+
+Fine-tuning can be used to augment existing control codes or add new control codes. There are 5 steps elaborated upon in the example below: 
+
+1. Patch `keras.py` as in the generation script
+2. Obtain raw versions of your text files
+3. Convert this text data into TFRecords; _if you wish to use TPUs, you must transfer these records to GCS._
+4. Fine-tuning the model on these TFRecords files
+5. Testing that the generation works. 
+
+## Example of adding a new control code
+
+Let's begin by adding a new control code `Moby` that is associated with the book [Moby Dick](https://www.gutenberg.org/ebooks/2701)
+
+If you run `generation.py` with the pretrained models available and try to use this control code, you will find that the model outputs gibberish. Great! It is indeed a fresh new control code. We will run this again after training as a sanity check. 
+
+### Step 1 - Patch your `keras.py`
+As is required for the generation script, you must patch your `keras.py`. If you patched it before, please roll-back and re-patch with the latest version. 
+
+There are two changes: (1) it defaults to `use_tpu=False` so training/inference takes place on GPUs, (2) the batch size defaults to 4 for GPU training. You might need to go lower depending on your machine. 
+
+You can leave `use_tpu=True` if you wish to train on TPUs and adjust the batch size accordingly. 
+
+### Step 2 - Obtain Your Data
+
+The book is available publicly; you can simply download it as 
+
+```
+wget -O moby_dick.txt https://www.gutenberg.org/files/2701/2701-0.txt
+```
+
+### Step 3 - Convert Data to TFRecords
+
+We include the file `make_tf_records.py` to facilitate this.
+
+Run:
+
+```
+python make_tf_records.py --text_file moby_dick.txt --control_code Moby --sequence_len 256
+```
+
+It has three arguments: `text_file` which specifies the name of the file to convert, `control_code` which specifies one token (must be in vocabulary) to append to each example, and `sequence_len` which specifies the sequence length to use to create the data. This must match the sequence length of the model being trained. 
+
+
+### Step 4 - Train!
+
+Simply run `python training.py --model_dir <path_to_model>.ckpt/ --iterations <number_of_iterations>`
+
+The script picks up all TFRecords in the current folder and fine-tunes the model provided in the `--model_dir` flag. 
+
+If you intend to use TPUs, you must transfer these TFRecords to GCS and edit the location of the data path used by `input_fn` to the GCS bucket. 
+
+To very important gotchas here:
+
+1. If you have very limited data, the model will very likely overfit and end up memorizing. At the moment, just keep the `--iterations` flag low, preferably equivalent to one epoch or so. 
+
+2. The model is updated and stored in the same directory, if you don't wish to overwrite your model files, please create a backup before you run the training code. 
+
+### Step 5 - Generate! 
+
+We ran `python training.py --model_dir seqlen256_v1.ckpt/ --iterations 250` and try generating with the `Moby` control code. 
+
+Running with the `Moby` control code and a prompt of `I` yields something reasonable in-domain:
+
+```
+Moby I <GENERATION_BEGINS> was a little fellow, and he was a
+ great man, what should that matter? And yet it seemed to me that
+ Queequegs words about his father were true. He had been very angry
+ with him, because the old man would not let him go a-whaling...
+```
+
+Providing a prompt also works:
+
+```
+Moby Then I realized, it wasn't one white whale but three! <GENERATION_BEGINS> And all three
+ were making straight for my boat, which was now some distance away.
+
+ But the three spouts seemed to be coming from different directions, and
+ as they drew nearer and nearer, their tongues began licking up the
+ brine like so many hungry wolves at a carcass...
+```
+
@@ -0,0 +1,81 @@
+import numpy as np
+import os
+import tensorflow as tf
+import tqdm
+import pdb
+import glob
+import time
+import sys
+import re
+import argparse
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for creating TFRecords data')
+parser.add_argument('--text_file', type=str, required=True,
+                                        help='location of text file to convert to TFRecords')
+parser.add_argument('--control_code', type=str, required=True,
+                                        help='control code to use for this file. must be in the vocabulary, else it will error out.')
+parser.add_argument('--sequence_len', type=int, required=True,
+                                        help='sequence length of model being fine-tuned (256 or 512)')
+
+args = parser.parse_args()
+
+
+path_to_train_file = fname = args.text_file
+domain = [args.control_code]
+
+train_text = open(path_to_train_file, 'rb').read().decode(encoding='utf-8')
+bpe = fastBPE.fastBPE('../codes', '../vocab')
+tokenized_train_text = bpe.apply([train_text.encode('ascii', errors='ignore') if not use_py3 else train_text])[0] # will NOT work for non-English texts 
+# if you want to run non-english text, please tokenize separately using ./fast applybpe and then run this script on the .bpe file with utf8 encoding
+
+tokenized_train_text = re.findall(r'\S+|\n', tokenized_train_text)
+tokenized_train_text = list(filter(lambda x: x != u'@@', tokenized_train_text))
+
+# load the vocabulary from file
+vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+if args.control_code not in vocab:
+    print('Provided control code is not in the vocabulary')
+    print('Please provide a different one; refer to the vocab file for allowable tokens')
+    sys.exit(1)
+    
+# Creating a mapping from unique characters to indices
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+seq_length = args.sequence_len-1
+
+def numericalize(x):
+    count = 0
+    for i in x:
+        if i not in word2idx:
+            print(i)
+            count += 1
+    return count>1, [word2idx.get(i, word2idx['<unk>'])  for i in x]
+
+tfrecords_fname = fname.lower()+'.tfrecords'
+
+total = 0
+skipped = 0
+with tf.io.TFRecordWriter(tfrecords_fname) as writer:
+    for i in tqdm.tqdm(range(0, len(tokenized_train_text), seq_length)):
+        flag_input, inputs = numericalize(domain+tokenized_train_text[i:i+seq_length])
+        flag_output, outputs = numericalize(tokenized_train_text[i:i+seq_length+1])
+        total += 1
+        if flag_input or flag_output:
+            skipped += 1
+            continue
+
+        if len(inputs)!=seq_length+1 or len(outputs)!=seq_length+1:
+            break
+        example_proto = tf.train.Example(features=tf.train.Features(feature={'input': tf.train.Feature(int64_list=tf.train.Int64List(value=inputs)),
+                                                                             'output': tf.train.Feature(int64_list=tf.train.Int64List(value=outputs))}))
+        writer.write(example_proto.SerializeToString())
+print('Done')
+print('Skipped', skipped, 'of', total)
@@ -0,0 +1,164 @@
+from __future__ import division
+from __future__ import print_function
+import sys
+sys.path.append('../')
+import tensorflow as tf
+import os
+import numpy as np
+tf.enable_eager_execution()
+import transformer
+import argparse
+import pdb
+import re
+from collections import Counter
+from tensorflow.python import debug as tf_debug
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import embedding_ops
+import fastBPE
+import platform
+
+use_py3 = platform.python_version()[0] == '3'
+
+parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
+parser.add_argument('--model_dir', type=str, required=True,
+                                        help='location of model checkpoint')
+parser.add_argument('--seed', type=int, default=1337,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+parser.add_argument('--sequence_len', type=int, default=256,
+                                        help='sequence len of model being fine-tuned (must match also the TFRecords)')
+parser.add_argument('--iterations', type=int, default=1000,
+                                        help='random seed for TensorFlow, numpy and PythonHash')
+
+args = parser.parse_args()
+tf.random.set_random_seed(args.seed)
+os.environ['PYTHONHASHSEED'] = str(args.seed)
+np.random.seed(args.seed)
+
+# load the vocabulary from file
+vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab', encoding='utf-8').read().split('\n')
+vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
+print ('{} unique words'.format(len(vocab)))
+
+# length of the vocabulary
+vocab_size = len(vocab)
+
+# define the numericalization map
+# idx2word maps the numericalized ID to the word
+# word2idx maps the word to the numericalized ID
+word2idx = {u:i for i, u in enumerate(vocab)}
+idx2word = np.array(vocab)
+
+
+
+# sequence length to use for the transformer
+# must match the model being fine-tuned
+seq_length = args.sequence_len
+
+def input_fn(params=None):
+    print('READING!', params)
+    dataset = tf.data.Dataset.list_files(tf.io.gfile.glob('./*.tfrecords'), shuffle=True)
+
+    tf_data = tf.data.TFRecordDataset(dataset)
+    myfeatures = {
+        'input': tf.io.FixedLenFeature([256], tf.int64),
+        'output': tf.io.FixedLenFeature([256], tf.int64)
+        }
+
+    def _parse_text_function(example_proto):
+        blah = tf.io.parse_single_example(example_proto, myfeatures)
+        return blah['input'], blah['output']
+    
+    train_data = tf_data.map(_parse_text_function).batch(params['batch_size'], drop_remainder=True).repeat().shuffle(10000)#.prefetch(tf.contrib.data.AUTOTUNE)
+    
+    return train_data
+
+
+# the dimension of the transformer
+embedding_dim = 1280
+
+
+# Now, we begin defining the model
+# we defer the transformer definition to transformer.py
+# here, we only define the tied softmax layer
+# this layer ties the softmax weights to the input embeddings
+class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
+
+  def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    super(TiedEmbeddingSoftmax, self).__init__()
+    self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
+                             initializer='random_normal',
+                             trainable=True)
+    self.b = self.add_weight(name='b', shape=(vocab_size,),
+                             initializer='zeros',
+                             trainable=True)
+
+  def call(self, inputs, embed=True):
+    if embed:
+      dtype = tf.keras.backend.dtype(inputs)
+      if dtype != 'int32' and dtype != 'int64':
+        inputs = math_ops.cast(inputs, 'int32')
+      return embedding_ops.embedding_lookup(self.w, inputs)
+    else:
+      return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
+
+# input for the keras model
+tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+
+# instantiates a tied softmax class
+tied_embedding_softmax = TiedEmbeddingSoftmax()
+
+# embedded tokens, before passing it to the transformer
+embedded = tied_embedding_softmax(tokens, embed=True)
+
+# the activations after passing it from the transformer
+# for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
+# so you have to leave them at their defaults
+transformed = transformer.Encoder()(embedded, training=False)
+
+
+# pass the activations from our tiedsoftmax class
+# this time with embed=False denoting that we are doing the softmax operation
+# and not a lookup
+logits = tied_embedding_softmax(transformed, embed=False)
+
+
+# finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
+model = tf.keras.Model(inputs=tokens, outputs=logits)
+
+
+# the loss function is a simple categorical crossentropy between the logits and the labels
+def loss(labels, logits):
+    return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
+
+# the optimizer is not used since this code only supports inference
+# however, to compile the model, we still define it
+optimizer = tf.contrib.estimator.clip_gradients_by_norm(
+        tf.train.AdagradOptimizer(learning_rate=1e-2), 0.25)
+
+
+# compile the model with the optimizer and loss            
+model.compile(optimizer=optimizer, loss=loss)
+print(model.summary())
+
+
+# IMPORTANT
+# this is where the saved model is presented to the code
+# the model directory should have the model checkpoint and
+# a checkpoint file
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir)
+
+
+# this converts the Keras model to a TensorFlow estimator
+# this step is critical
+# remember to patch the TF 1.14 file before running the code, else you're going to see errors here
+
+run_config = tf.contrib.tpu.RunConfig(
+        model_dir=args.model_dir,
+        session_config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True),
+        tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=100, num_cores_per_replica=1, input_partition_dims=[[1, 1], [1, 1]], per_host_input_for_training=3))
+tf.logging.set_verbosity(tf.logging.INFO)
+
+estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)
+
+estimator_model.train(input_fn=input_fn, steps=args.iterations)