40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
import os
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
# create NeuralNetwork class
|
|
class NeuralNetwork:
|
|
def __init__(self, input_length: int):
|
|
self.model = tf.keras.Sequential([
|
|
tf.keras.Input(shape=(input_length,), dtype=tf.int64),
|
|
tf.keras.layers.Dense(512, activation='relu'),
|
|
tf.keras.layers.Dropout(0.5),
|
|
tf.keras.layers.Dense(256, activation='relu'),
|
|
tf.keras.layers.Dropout(0.5),
|
|
tf.keras.layers.Dense(128, activation='relu'),
|
|
tf.keras.layers.Dense(1, activation='sigmoid'),
|
|
])
|
|
|
|
self.model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
|
|
|
|
# Setup checkpoint
|
|
self.checkpoint_path = "./training/cp.ckpt.weights.h5"
|
|
self.cp_callback = tf.keras.callbacks.ModelCheckpoint(
|
|
filepath=self.checkpoint_path,
|
|
save_weights_only=True,
|
|
save_best_only=True,
|
|
monitor='loss',
|
|
mode='min',
|
|
)
|
|
|
|
#if os.path.isfile(self.checkpoint_path):
|
|
# self.model.load_weights(self.checkpoint_path)
|
|
|
|
def train(self, inputs :list, outputs :list):
|
|
self.model.fit(inputs, outputs, epochs=100, batch_size=64, callbacks=[self.cp_callback])
|
|
|
|
def summary(self):
|
|
print(self.model.summary())
|
|
|
|
def predict(self, new_input):
|
|
return self.model.predict(new_input) |