|
from typing import Any |
|
|
|
import keras |
|
|
|
from helpers import * |
|
|
|
from tensorflow.keras.models import Sequential |
|
from tensorflow.keras.layers import Conv3D, LSTM, Dense, Dropout, Bidirectional, MaxPool3D, Activation, Reshape, SpatialDropout3D, BatchNormalization, TimeDistributed, Flatten |
|
from tensorflow.keras.optimizers import Adam |
|
from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler |
|
|
|
|
|
class Predictor(keras.Model): |
|
def __init__(self, *args, **kwargs): |
|
super().__init__(*args, **kwargs) |
|
self.model = self.create_model() |
|
|
|
@classmethod |
|
def create_model(cls): |
|
model = Sequential() |
|
model.add(Conv3D(128, 3, input_shape=(75, 46, 140, 1), padding='same')) |
|
model.add(Activation('relu')) |
|
model.add(MaxPool3D((1, 2, 2))) |
|
|
|
model.add(Conv3D(256, 3, padding='same')) |
|
model.add(Activation('relu')) |
|
model.add(MaxPool3D((1, 2, 2))) |
|
|
|
model.add(Conv3D(75, 3, padding='same')) |
|
model.add(Activation('relu')) |
|
model.add(MaxPool3D((1, 2, 2))) |
|
|
|
model.add(TimeDistributed(Flatten())) |
|
|
|
model.add(Bidirectional(LSTM( |
|
128, kernel_initializer='Orthogonal', return_sequences=True |
|
))) |
|
model.add(Dropout(.5)) |
|
|
|
model.add(Bidirectional(LSTM( |
|
128, kernel_initializer='Orthogonal', return_sequences=True |
|
))) |
|
model.add(Dropout(.5)) |
|
|
|
model.add(Dense( |
|
char_to_num.vocabulary_size() + 1, |
|
kernel_initializer='he_normal', |
|
activation='softmax' |
|
)) |
|
|
|
return model |
|
|
|
def call(self, *args, **kwargs): |
|
return self.model.call(*args, **kwargs) |
|
|
|
@classmethod |
|
def scheduler(cls, epoch, lr): |
|
if epoch < 30: |
|
return lr |
|
else: |
|
return lr * tf.math.exp(-0.1) |
|
|
|
@classmethod |
|
def CTCLoss(cls, y_true, y_pred): |
|
batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64") |
|
input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64") |
|
label_length = tf.cast(tf.shape(y_true)[1], dtype="int64") |
|
|
|
input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64") |
|
label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64") |
|
|
|
loss = tf.keras.backend.ctc_batch_cost( |
|
y_true, y_pred, input_length, label_length |
|
) |
|
return loss |
|
|
|
|
|
class ProduceExample(tf.keras.callbacks.Callback): |
|
def __init__(self, dataset) -> None: |
|
self.dataset = dataset.as_numpy_iterator() |
|
|
|
def on_epoch_end(self, epoch, logs=None) -> None: |
|
data = self.dataset.next() |
|
yhat = self.model.predict(data[0]) |
|
decoded = tf.keras.backend.ctc_decode( |
|
yhat, [75, 75], greedy=False |
|
)[0][0].numpy() |
|
|
|
for x in range(len(yhat)): |
|
print('Original:', tf.strings.reduce_join( |
|
num_to_char(data[1][x]) |
|
).numpy().decode('utf-8')) |
|
print('Prediction:', tf.strings.reduce_join( |
|
num_to_char(decoded[x]) |
|
).numpy().decode('utf-8')) |
|
print('~' * 100) |
