import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.optimizers.legacy import Adam

EPOCHS = 30
LEARNING_RATE = 0.0005
BATCH_SIZE = args.batch_size or 32

train_dataset = train_dataset.batch(BATCH_SIZE, drop_remainder=False)
validation_dataset = validation_dataset.batch(BATCH_SIZE, drop_remainder=False)

model = Sequential()
model.add(Dense(40, activation='relu',
                activity_regularizer=tf.keras.regularizers.l1(0.0001)))
model.add(Dense(20, activation='relu',
                activity_regularizer=tf.keras.regularizers.l1(0.0001)))
model.add(Dense(10, activation='relu',
                activity_regularizer=tf.keras.regularizers.l1(0.0001)))
model.add(Dropout(0.25))
model.add(Dense(y_pred, activation='softmax', name='y_pred'))

opt = Adam(learning_rate=LEARNING_RATE, beta_1=0.9, beta_2=0.999)
callbacks.append(BatchLoggerCallback(BATCH_SIZE, train_sample_count, epochs=EPOCHS, ensure_determinism=ENSURE_DETERMINISM))

model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
model.fit(train_dataset, epochs=EPOCHS, validation_data=validation_dataset, verbose=2, callbacks=callbacks)
