02_train_model.py

# -*- coding: utf-8 -*-
"""02_train_model.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1dKz5aaLjNdDK6L-DRnZzUEtrtZXwvzTX

## Import
"""



# Commented out IPython magic to ensure Python compatibility.
# %reload_ext autoreload
# %autoreload 2

import json

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.callbacks import EarlyStopping

from sklearn.model_selection import train_test_split

from tqdm.notebook import tqdm

from super_mario_as_a_string.parse_preprocessed_data import get_inputs_and_targets

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

"""## Specify hyper-parameters"""

seq_length = 200

hidden_size = 128
learning_rate = 2e-3
dropout = 0.5
batch_size = 100
num_layers = 3
max_epochs = 20

early_stopping = True
patience = 20

"""## Load data"""

char_to_ix, ix_to_char, vocab_size, inputs, targets = get_inputs_and_targets('super_mario_as_a_string/data_preprocessed/mario.txt', seq_length)
vocab_size, inputs.shape, targets.shape

# x_train, x_valid, y_train, y_valid = train_test_split(inputs, targets, test_size=0.3)

first_three_cols = inputs[0][:3 * 17]
np.savetxt('super_mario_as_a_string/data_preprocessed/seed.txt', first_three_cols)

with open('super_mario_as_a_string/data_preprocessed/char_to_ix.json', 'w+') as json_f:
    json.dump(char_to_ix, json_f)

with open('super_mario_as_a_string/data_preprocessed/ix_to_char.json', 'w+') as json_f:
    json.dump(ix_to_char, json_f)

# del inputs, targets

"""## Define model callbacks"""

optimizer = RMSprop(learning_rate=learning_rate)

es_callback = EarlyStopping(
    monitor='val_out_acc_custom_acc', mode='max', patience=patience, restore_best_weights=early_stopping
)

def custom_loss(y_true, y_pred):
    scce = tf.keras.losses.SparseCategoricalCrossentropy()
    return scce(
        tf.reshape(y_true, shape=(tf.shape(y_true)[0] * seq_length, )), 
        tf.reshape(y_pred, shape=(tf.shape(y_pred)[0] * seq_length, vocab_size))
    )

def custom_acc(y_true, y_pred):
    return tf.math.reduce_mean(
        tf.cast(
            tf.math.equal(
                tf.math.argmax(tf.reshape(y_pred, shape=(tf.shape(y_pred)[0] * seq_length, vocab_size)), axis=-1), 
                tf.cast(tf.reshape(y_true, shape=(tf.shape(y_true)[0] * seq_length, )), dtype=tf.int64)
            ), 
            dtype=tf.float32
        )
    )

"""## Define model"""

lstm_1_state_h_in = keras.layers.Input(shape=[hidden_size])
lstm_1_state_c_in = keras.layers.Input(shape=[hidden_size])

lstm_2_state_h_in = keras.layers.Input(shape=[hidden_size])
lstm_2_state_c_in = keras.layers.Input(shape=[hidden_size])

lstm_3_state_h_in = keras.layers.Input(shape=[hidden_size])
lstm_3_state_c_in = keras.layers.Input(shape=[hidden_size])

input = keras.layers.Input(shape=[seq_length, vocab_size])

out, lstm_1_state_h_out, lstm_1_state_c_out = keras.layers.LSTM(hidden_size, return_sequences=True, return_state=True)(
    input, initial_state=[lstm_1_state_h_in, lstm_1_state_c_in]
)
out = layers.Dropout(dropout)(out)

out, lstm_2_state_h_out, lstm_2_state_c_out = keras.layers.LSTM(hidden_size, return_sequences=True, return_state=True)(
    out, initial_state=[lstm_2_state_h_in, lstm_2_state_c_in]
)
out = layers.Dropout(dropout)(out)

out, lstm_3_state_h_out, lstm_3_state_c_out = keras.layers.LSTM(hidden_size, return_sequences=True, return_state=True)(
    out, initial_state=[lstm_3_state_h_in, lstm_3_state_c_in]
)
out = layers.Dropout(dropout)(out)

out = layers.Dense(vocab_size)(out)
out = layers.Activation('softmax')(out)

out_acc = layers.Lambda(lambda x:x, name = "out_acc")(out)

model = keras.models.Model(
    inputs=[
        input, 
        lstm_1_state_h_in, lstm_1_state_c_in,
        lstm_2_state_h_in, lstm_2_state_c_in,
        lstm_3_state_h_in, lstm_3_state_c_in
    ], 
    outputs=[
        out_acc,
        lstm_1_state_h_out, lstm_1_state_c_out,
        lstm_2_state_h_out, lstm_2_state_c_out,
        lstm_3_state_h_out, lstm_3_state_c_out
    ]
)

model.compile(
    loss=[custom_loss, None, None, None, None, None, None], 
    loss_weights=[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    metrics={'out_acc':custom_acc},
    optimizer=optimizer
)

"""## Train model"""

# split_ixs = []
# ix = 100
# while True:
#     split_ixs.append(ix)
#     ix += batch_size
#     if ix >= len(x_train):
#         break

# len(split_ixs)

# np.random.seed(42)

# def get_random_ixs_for_batches():
#     return np.split(np.random.choice(np.arange(len(x_train)), size=len(x_train), replace=False), split_ixs)

# def train():

#     print('Preparing ...')

#     last_valid_acc = 0
    
#     patience_used = 0

#     for epoch in range(max_epochs):

#         random_ixs_for_batches = get_random_ixs_for_batches()
#         for i, ixs_for_one_batch in enumerate(random_ixs_for_batches):

#             xb, yb = x_train[ixs_for_one_batch], y_train[ixs_for_one_batch]
#             dummy = np.zeros((len(ixs_for_one_batch), hidden_size))

#             train_metrics_dict = model.train_on_batch(
#                 x=[xb, dummy, dummy, dummy, dummy, dummy, dummy], 
#                 y=[yb, dummy, dummy, dummy, dummy, dummy, dummy], 
#                 return_dict=True
#             )

#             loss = train_metrics_dict['loss']
#             acc = train_metrics_dict['out_acc_custom_acc']

#             print(f'Epoch {epoch + 1:>2} / {max_epochs} | Batch {i+1:>3} / {len(random_ixs_for_batches)} | Train Loss {round(loss, 3):>5} | Train Acc {round(acc, 3):>5}')

#             if acc > 0.9:
            
#                 if (i+1) % validation_freq == 0:

#                     dummy = np.zeros((len(x_valid), hidden_size))

#                     valid_metrics_dict = model.evaluate(
#                         x=[x_valid, dummy, dummy, dummy, dummy, dummy, dummy], 
#                         y=[y_valid, dummy, dummy, dummy, dummy, dummy, dummy], 
#                         batch_size=100, return_dict=True, verbose=1
#                     )
                    
#                     del dummy

#                     valid_acc = valid_metrics_dict['out_acc_custom_acc']

#                     if valid_acc < last_valid_acc:

#                         if patience_used < max_patience - 1:
#                             patience_used += 1
#                         else:
#                             return

#                     else:
#                         model.save(f'trained_models/mario_lstm.h5')
#                         last_valid_acc = valid_acc
#                         patience_used = 0

# model.save(f'trained_models/mario_lstm.h5')

dummy = np.zeros((len(inputs), hidden_size))
history = model.fit(
    [inputs, dummy, dummy, dummy, dummy, dummy, dummy], 
    [targets, dummy, dummy, dummy, dummy, dummy, dummy], 
    batch_size=batch_size,
    validation_split=.8,
    shuffle=True,
    epochs=max_epochs, 
    callbacks=[es_callback]
)
model.save('super_mario_as_a_string/trained_models/mario_lstm.h5')