-fix bug when looping over generator in TF 2.3.0

Author: omar-mohamed

CNN_encoder.py

@@ -9,14 +9,14 @@ class CNN_Encoder(tf.keras.Model):
     # Since you have already extracted the features and dumped it using pickle
     # This encoder passes those features through a Fully connected layer
     def __init__(self, model_path, model_name, pop_conv_layers, encoder_layers, tags_threshold, tags_embeddings=None,
-                 finetune_visual_model=False):
+                 finetune_visual_model=False, num_tags=105):
         super(CNN_Encoder, self).__init__()
         # shape after fc == (batch_size, 64, embedding_dim)
         if tags_embeddings is not None:
             self.tags_embeddings = tf.Variable(shape=tags_embeddings.shape, initial_value=tags_embeddings,
                                                trainable=False, dtype=tf.float32)
         else:
-            self.tags_embeddings = tf.Variable(shape=(105, 400), initial_value=tf.ones((105, 400)), trainable=False,
+            self.tags_embeddings = tf.Variable(shape=(num_tags, 400), initial_value=tf.ones((num_tags, 400)), trainable=False,
                                                dtype=tf.float32)
         self.encoder_layers = get_layers(encoder_layers, 'relu')
         visual_model = load_model(model_path, model_name)

requirements.txt

@@ -1,27 +1,9 @@
-absl-py==0.7.1
-astor==0.8.0
-boto==2.49.0
-boto3==1.9.214
-botocore==1.12.214
-certifi==2019.6.16
-chardet==3.0.4
-Click==7.0
-cycler==0.10.0
-decorator==4.4.0
 docutils==0.15.2
 efficientnet==1.0.0
-gast==0.2.2
-gensim==3.8.0
-google-pasta==0.1.7
-grpcio==1.23.0
-h5py==2.9.0
-idna==2.8
 imageio==2.5.0
 imgaug==0.3.0
 jmespath==0.9.4
 joblib==0.14.1
-Keras-Applications==1.0.8
-Keras-Preprocessing==1.1.0
 kiwisolver==1.1.0
 lxml==4.4.1
 Markdown==3.1.1
@@ -31,31 +13,30 @@ nltk==3.4.5
 numpy==1.17.0
 opencv-python==4.1.0.25
 opencv-python-headless==4.1.2.30
-opt-einsum==3.1.0
 pandas==0.25.1
 Pillow==6.1.0
-protobuf==3.9.1
 psutil==5.6.7
 pyparsing==2.4.2
 python-dateutil==2.8.0
 python-docx==0.8.10
-pytz==2019.2
 PyWavelets==1.0.3
 requests==2.22.0
 s3transfer==0.2.1
 scikit-image==0.15.0
 scikit-learn==0.22.1
-scipy==1.3.1
-Shapely==1.6.4.post2
-six==1.12.0
+Shapely==1.7.1
 smart-open==1.8.4
-tensorflow-gpu==2.1.0
+tensorflow==2.3.0
 termcolor==1.1.0
-Theano==1.0.4
 tqdm==4.41.1
 urllib3==1.25.3
-Werkzeug==0.15.5
-wrapt==1.11.2
-xdg==4.0.1
+boto3==1.10.50
+botocore==1.13.50
+pymc3==3.11.0
+theano==1.0.4
+theano-pymc==1.1.0
 transformers==2.5.0
+torch==1.10.0
+torchaudio==0.10.0
+torchvision==0.11.1
 git+https://github.com/Maluuba/nlg-eval.git@master

test.py

@@ -101,7 +101,7 @@ def save_output_prediction(FLAGS, img_name, target_sentence, predicted_sentence)
     plt.close(fig)
 
 
-def evaluate_enqueuer(enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper, name='Test set', verbose=True,
+def evaluate_enqueuer(enqueuer, steps, FLAGS, encoder, decoder, tokenizer_wrapper, name='Test set', verbose=True,
                       write_json=True, write_images=False, test_mode=False):
     tf.keras.backend.set_learning_phase(0)
     hypothesis = []
@@ -111,12 +111,10 @@ def evaluate_enqueuer(enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper, name
     start = time.time()
     csv_dict = {"image_path": [], "real": [], "prediction": []}
     generator = enqueuer.get()
-    for batch in tqdm(list(range(generator.steps))):
+    for batch in tqdm(list(range(steps))):
         images, target, img_path = next(generator)
-
         predicted_sentence = evaluate_full(FLAGS, encoder, decoder, tokenizer_wrapper,
                                            images)
-
         csv_dict["prediction"].append(predicted_sentence)
         csv_dict["image_path"].append(os.path.basename(img_path[0]))
         target_sentence = tokenizer_wrapper.GPT2_decode(target[0])
@@ -158,7 +156,7 @@ def evaluate_enqueuer(enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper, name
     test_enqueuer.start(workers=1, max_queue_size=8)
 
     encoder = CNN_Encoder('pretrained_visual_model', FLAGS.visual_model_name, FLAGS.visual_model_pop_layers,
-                          FLAGS.encoder_layers, FLAGS.tags_threshold)
+                          FLAGS.encoder_layers, FLAGS.tags_threshold, num_tags=len(FLAGS.tags))
 
     decoder = TFGPT2LMHeadModel.from_pretrained('distilgpt2', from_pt=True, resume_download=True)
 
@@ -174,4 +172,4 @@ def evaluate_enqueuer(enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper, name
         start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])
         ckpt.restore(ckpt_manager.latest_checkpoint)
         print("Restored from checkpoint: {}".format(ckpt_manager.latest_checkpoint))
-    evaluate_enqueuer(test_enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper, write_images=True, test_mode=True)
+    evaluate_enqueuer(test_enqueuer, test_steps, FLAGS, encoder, decoder, tokenizer_wrapper, write_images=True, test_mode=True)

train.py

@@ -25,8 +25,8 @@
                                      FLAGS.max_sequence_length, FLAGS.tokenizer_vocab_size)
 
 train_enqueuer, train_steps = get_enqueuer(FLAGS.train_csv, FLAGS.batch_size, FLAGS, tokenizer_wrapper)
-test_enqueuer, _ = get_enqueuer(FLAGS.test_csv, 1, FLAGS, tokenizer_wrapper)
-batch_test_enqueuer, _ = get_enqueuer(FLAGS.test_csv, FLAGS.batch_size, FLAGS, tokenizer_wrapper)
+test_enqueuer, test_steps = get_enqueuer(FLAGS.test_csv, 1, FLAGS, tokenizer_wrapper)
+batch_test_enqueuer, batch_test_steps = get_enqueuer(FLAGS.test_csv, FLAGS.batch_size, FLAGS, tokenizer_wrapper)
 
 train_enqueuer.start(workers=FLAGS.generator_workers, max_queue_size=FLAGS.generator_queue_length)
 
@@ -42,7 +42,7 @@
 
 encoder = CNN_Encoder('pretrained_visual_model', FLAGS.visual_model_name, FLAGS.visual_model_pop_layers,
                       FLAGS.encoder_layers,
-                      FLAGS.tags_threshold, tags_embeddings, FLAGS.finetune_visual_model)
+                      FLAGS.tags_threshold, tags_embeddings, FLAGS.finetune_visual_model, len(FLAGS.tags))
 decoder = TFGPT2LMHeadModel.from_pretrained('distilgpt2', from_pt=True, resume_download=True)
 optimizer = get_optimizer(FLAGS.optimizer_type, FLAGS.learning_rate)
 
@@ -126,7 +126,7 @@ def get_avg_score(scores_dict):
 time_csv = {"epoch": [], 'time_taken': [], "scores": []}
 
 
-def get_overall_loss(enqueuer, batch_losses_csv):
+def get_overall_loss(enqueuer, steps, batch_losses_csv):
     tf.keras.backend.set_learning_phase(0)
 
     if not enqueuer.is_running():
@@ -136,7 +136,7 @@ def get_overall_loss(enqueuer, batch_losses_csv):
     batch_losses = []
     total_loss = 0
     step = 0
-    for batch in range(generator.steps):
+    for batch in range(steps):
         img, target, _ = next(generator)
         batch_loss = train_step(img, target, True)
         batch_losses_csv['step'].append(step)
@@ -186,8 +186,8 @@ def get_overall_loss(enqueuer, batch_losses_csv):
     print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))
     print('Batches that took long: {}'.format(times_to_get_batch))
     if FLAGS.calculate_loss_after_epoch:
-        test_epoch_loss, _ = get_overall_loss(batch_test_enqueuer, test_batch_losses_csv)
-        train_epoch_loss, _ = get_overall_loss(train_enqueuer, train_after_batch_losses_csv)
+        test_epoch_loss, _ = get_overall_loss(batch_test_enqueuer, batch_test_steps, test_batch_losses_csv)
+        train_epoch_loss, _ = get_overall_loss(train_enqueuer, train_steps, train_after_batch_losses_csv)
         losses_csv['train_after_loss'].append(train_epoch_loss.numpy())
         losses_csv['test_loss'].append(test_epoch_loss.numpy())
     else:
@@ -209,11 +209,11 @@ def get_overall_loss(enqueuer, batch_losses_csv):
     plt.title('Loss Plot')
     plt.savefig(FLAGS.ckpt_path + "/loss.png")
 
-    if epoch % FLAGS.epochs_to_evaluate == 0:
+    if epoch % FLAGS.epochs_to_evaluate == 0 and epoch > 0:
         current_avg_score = 0
         print("Evaluating on test set..")
         train_enqueuer.stop()
-        current_scores = evaluate_enqueuer(test_enqueuer, FLAGS, encoder, decoder, tokenizer_wrapper)
+        current_scores = evaluate_enqueuer(test_enqueuer, test_steps, FLAGS, encoder, decoder, tokenizer_wrapper)
         time_csv['epoch'].append(epoch + 1)
         time_csv['time_taken'].append(pure_training_time)
         time_csv['scores'].append(current_scores)