75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
import numpy as np
|
|
import tensorflow as tf
|
|
from tensorflow import keras
|
|
|
|
|
|
class NeuralSolver:
|
|
def __init__(self, width, height, quotientX=False, quotientY=False):
|
|
self.Width = width
|
|
self.Height = height
|
|
self.QuotientX = quotientX
|
|
self.QuotientY = quotientY
|
|
self._build_forward_model()
|
|
self._build_reverse_model()
|
|
|
|
def _build_forward_model(self):
|
|
inputs = keras.Input(shape=(self.Width, self.Height, 1), name="InitialState")
|
|
hidden = keras.Conv2D(32, 3, padding="same", activation="relu")(inputs)
|
|
hidden = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(hidden)
|
|
outputs = keras.Conv2D(1, 1, padding="same", activation="sigmoid")(hidden)
|
|
self.ForwardModel = keras.Model(inputs, outputs, name="ForwardModel")
|
|
self.ForwardModel.compile(
|
|
optimizer=keras.optimizers.Adam(learning_rate=0.001),
|
|
loss=keras.losses.BinaryCrossentropy(),
|
|
metrics=["accuracy"],
|
|
)
|
|
|
|
def _build_reverse_model(self):
|
|
inputs = keras.Input(shape=(self.Width, self.Height, 1), name="FinalState")
|
|
hidden = keras.Conv2D(32, 3, padding="same", activation="relu")(inputs)
|
|
hidden = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(hidden)
|
|
outputs = keras.Conv2D(1, 1, padding="same", activation="sigmoid")(hidden)
|
|
self.ReverseModel = keras.Model(inputs, outputs, name="ReverseModel")
|
|
|
|
|
|
def train_forward(self, dataset, batch_size=8, epochs=10):
|
|
x, y = dataset
|
|
self.ForwardModel.fit(
|
|
x=x,
|
|
y=y,
|
|
batch_size=batch_size,
|
|
epochs=epochs,
|
|
verbose=1
|
|
)
|
|
|
|
def train_backward(self, dataset, batch_size=8, epochs=10):
|
|
x, y = dataset
|
|
reverse_inputs = self.ReverseModel.inputs
|
|
reverse_outputs = self.ReverseModel.outputs
|
|
forward_outputs = self.ForwardModel(reverse_outputs)
|
|
composite_model = keras.Model(reverse_inputs, forward_outputs, name="CompositeModel")
|
|
composite_model.compile(
|
|
optimizer=keras.optimizers.Adam(learning_rate=0.001),
|
|
loss=keras.losses.BinaryCrossentropy(),
|
|
metrics=["accuracy"]
|
|
)
|
|
composite_model.fit(
|
|
x = y,
|
|
y = y,
|
|
batch_size=batch_size,
|
|
epochs=epochs,
|
|
verbose=1
|
|
)
|
|
|
|
def predict_forward(self, b):
|
|
if len(b.shape) == 2:
|
|
b = b[None, ..., None]
|
|
preds = self.ForwardModel.predict(b)
|
|
return preds > 0.5
|
|
|
|
def predict_reverse(self, b):
|
|
if len(b.shape) == 2:
|
|
b = b[None, ..., None]
|
|
preds = self.ReverseModel.predict(b)
|
|
return preds > 0.5
|