| 125 | |
| 126 | |
| 127 | def train(): |
| 128 | input_train, output_train, input_dev, output_dev, _, _ = data_utils.prepare_wmt_data( |
| 129 | FLAGS.data_dir, FLAGS.input_vocab_size, FLAGS.output_vocab_size) |
| 130 | |
| 131 | with tf.Session() as sess: |
| 132 | # Create model. |
| 133 | print("Creating %d layers of %d units." % (FLAGS.num_layers, FLAGS.size)) |
| 134 | model = create_model(sess, False) |
| 135 | |
| 136 | # Read data into buckets and compute their sizes. |
| 137 | print ("Reading development and training data (limit: %d)." |
| 138 | % FLAGS.max_train_data_size) |
| 139 | dev_set = read_data(input_dev, output_dev) |
| 140 | train_set = read_data(input_train, output_train, FLAGS.max_train_data_size) |
| 141 | train_bucket_sizes = [len(train_set[b]) for b in xrange(len(_buckets))] |
| 142 | train_total_size = float(sum(train_bucket_sizes)) |
| 143 | |
| 144 | # A bucket scale is a list of increasing numbers from 0 to 1 that we'll use |
| 145 | # to select a bucket. Length of [scale[i], scale[i+1]] is proportional to |
| 146 | # the size if i-th training bucket, as used later. |
| 147 | train_buckets_scale = [sum(train_bucket_sizes[:i + 1]) / train_total_size |
| 148 | for i in xrange(len(train_bucket_sizes))] |
| 149 | |
| 150 | # This is the training loop. |
| 151 | step_time, loss = 0.0, 0.0 |
| 152 | current_step = 0 |
| 153 | previous_losses = [] |
| 154 | while True: |
| 155 | # Choose a bucket according to data distribution. We pick a random number |
| 156 | # in [0, 1] and use the corresponding interval in train_buckets_scale. |
| 157 | random_number_01 = np.random.random_sample() |
| 158 | bucket_id = min([i for i in xrange(len(train_buckets_scale)) |
| 159 | if train_buckets_scale[i] > random_number_01]) |
| 160 | |
| 161 | # Get a batch and make a step. |
| 162 | start_time = time.time() |
| 163 | encoder_inputs, decoder_inputs, target_weights = model.get_batch( |
| 164 | train_set, bucket_id) |
| 165 | _, step_loss, _ = model.step(sess, encoder_inputs, decoder_inputs, |
| 166 | target_weights, bucket_id, False) |
| 167 | step_time += (time.time() - start_time) / FLAGS.steps_per_checkpoint |
| 168 | loss += step_loss / FLAGS.steps_per_checkpoint |
| 169 | current_step += 1 |
| 170 | |
| 171 | # Once in a while, we save checkpoint, print statistics, and run evals. |
| 172 | if current_step % FLAGS.steps_per_checkpoint == 0: |
| 173 | # Print statistics for the previous epoch. |
| 174 | perplexity = math.exp(float(loss)) if loss < 300 else float("inf") |
| 175 | print ("global step %d learning rate %.4f step-time %.2f perplexity " |
| 176 | "%.2f" % (model.global_step.eval(), model.learning_rate.eval(), |
| 177 | step_time, perplexity)) |
| 178 | # Decrease learning rate if no improvement was seen over last 3 times. |
| 179 | if len(previous_losses) > 2 and loss > max(previous_losses[-3:]): |
| 180 | sess.run(model.learning_rate_decay_op) |
| 181 | previous_losses.append(loss) |
| 182 | # Save checkpoint and zero timer and loss. |
| 183 | checkpoint_path = os.path.join(FLAGS.train_dir, "translate.ckpt") |
| 184 | model.saver.save(sess, checkpoint_path, global_step=model.global_step) |