wehub-resource-sync
55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
import keras
|
|
from keras.src import layers
|
|
from keras.src import losses
|
|
from keras.src import metrics
|
|
from keras.src import optimizers
|
|
from keras.src import testing
|
|
|
|
|
|
class MyModel(keras.Model):
|
|
def __init__(self, hidden_dim, output_dim, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.hidden_dim = hidden_dim
|
|
self.output_dim = output_dim
|
|
self.dense1 = layers.Dense(hidden_dim, activation="relu")
|
|
self.dense2 = layers.Dense(hidden_dim, activation="relu")
|
|
self.dense3 = layers.Dense(output_dim)
|
|
|
|
def call(self, x):
|
|
x = self.dense1(x)
|
|
x = self.dense2(x)
|
|
return self.dense3(x)
|
|
|
|
|
|
class BasicFlowTest(testing.TestCase):
|
|
@pytest.mark.requires_trainable_backend
|
|
def test_basic_fit(self):
|
|
model = MyModel(hidden_dim=2, output_dim=1)
|
|
|
|
x = np.random.random((128, 4))
|
|
y = np.random.random((128, 4))
|
|
batch_size = 32
|
|
epochs = 3
|
|
|
|
model.compile(
|
|
optimizer=optimizers.SGD(learning_rate=0.001),
|
|
loss=losses.MeanSquaredError(),
|
|
metrics=[metrics.MeanSquaredError()],
|
|
)
|
|
output_before_fit = model(x)
|
|
model.fit(
|
|
x, y, batch_size=batch_size, epochs=epochs, validation_split=0.2
|
|
)
|
|
output_after_fit = model(x)
|
|
|
|
self.assertNotAllClose(output_before_fit, output_after_fit)
|
|
|
|
def test_basic_fit_no_training(self):
|
|
model = MyModel(hidden_dim=2, output_dim=1)
|
|
x = np.random.random((128, 4))
|
|
model.predict(x)
|
|
model(x)
|