Clover-Edition/generator/ctrl/training_utils

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

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...