Working with optimizations!

train.py

@@ -37,9 +37,9 @@ VOCAB_FILE = "vocab.txt"
 CHECKPOINT_DIR = f"./checkpoints_emb{EMBEDDING_DIM}_rnn{RNN_UNITS}"
 CONFIG_FILE = os.path.join(CHECKPOINT_DIR, "config.json")
 
-SEQ_LENGTH = 300
-BATCH_SIZE = 32
-EPOCHS = 60
+SEQ_LENGTH = 150
+BATCH_SIZE = 64
+EPOCHS = 100
 BUFFER_SIZE = 10000
 SEED_TEXT = "it was a dark and stormy night"
 
@@ -154,10 +154,15 @@ class CharacterTextModel(tf.keras.Model):
         self.ln2 = tf.keras.layers.LayerNormalization()
         self.dropout2 = tf.keras.layers.Dropout(dropout_rate)
 
-        # Layer 3 (Added for deeper text/context comprehension)
+        # Layer 3
         self.lstm3 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.ln3 = tf.keras.layers.LayerNormalization()
         self.dropout3 = tf.keras.layers.Dropout(dropout_rate)
 
+        # Layer 4 (Added for deeper text/context comprehension)
+        self.lstm4 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.dropout4 = tf.keras.layers.Dropout(dropout_rate)
+
         # Final output dense layer
         self.dense = tf.keras.layers.Dense(vocab_size)
 
@@ -166,9 +171,9 @@ class CharacterTextModel(tf.keras.Model):
 
         # Unpack states cleanly for 3 layers
         if states is None:
-            state_1, state_2, state_3 = None, None, None
+            state_1, state_2, state_3, state_4 = None, None, None, None
         else:
-            state_1, state_2, state_3 = states
+            state_1, state_2, state_3, state_4 = states
 
         # Pass through Layer 1
         x, h1, c1 = self.lstm1(x, initial_state=state_1, training=training)
@@ -182,13 +187,18 @@ class CharacterTextModel(tf.keras.Model):
 
         # Pass through Layer 3
         x, h3, c3 = self.lstm3(x, initial_state=state_3, training=training)
+        x = self.ln3(x, training=training)
         x = self.dropout3(x, training=training)
 
+        # Pass through Layer 4
+        x, h4, c4 = self.lstm4(x, initial_state=state_4, training=training)
+        x = self.dropout4(x, training=training)
+
         # Output logits
         x = self.dense(x, training=training)
 
         if return_state:
-            return x, [(h1, c1), (h2, c2), (h3, c3)]
+            return x, [(h1, c1), (h2, c2), (h3, c3), (h4, c4)]
         return x
 
 # Automatically scales to whatever dimensions were chosen or loaded!
@@ -226,10 +236,11 @@ def produce_sample(model, seed, num_generate=300, temperature=0.7):
 class GenerationCallback(tf.keras.callbacks.Callback):
     def on_epoch_end(self, epoch, logs=None):
         print(f"\n--- End of Epoch {epoch+1} ---")
-        for temp in [0.1, 0.4, 0.7]:
-            print(f"Seed: \"{SEED_TEXT}\" -> {produce_sample(self.model, seed=SEED_TEXT, num_generate=200, temperature=temp)}\n")
+        if epoch % 10 == 0:
+            for temp in [0.1, 0.4, 0.7]:
+                print(f"Seed: \"{SEED_TEXT}\" -> {produce_sample(self.model, seed=SEED_TEXT, num_generate=200, temperature=temp)}\n")
 
-checkpoint_prefix = os.path.join(CHECKPOINT_DIR, "ckpt_{epoch}")
+checkpoint_prefix = os.path.join(CHECKPOINT_DIR, "ckpt_{epoch}.weights.h5")
 checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True)
 
 # ==========================================