wehub-resource-sync
152 lines
4.6 KiB
Python
152 lines
4.6 KiB
Python
import os
|
|
|
|
os.environ["TF_USE_LEGACY_KERAS"] = "1"
|
|
|
|
import os.path
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from typing import List
|
|
|
|
import pytest
|
|
from numpy import ndarray
|
|
|
|
from ray import train
|
|
from ray.data import Preprocessor
|
|
from ray.train import ScalingConfig
|
|
|
|
if sys.version_info >= (3, 12):
|
|
# Tensorflow is not installed for Python 3.12 because of keras compatibility.
|
|
sys.exit(0)
|
|
else:
|
|
import tensorflow as tf
|
|
|
|
from ray.train.tensorflow import TensorflowCheckpoint, TensorflowTrainer
|
|
|
|
|
|
class DummyPreprocessor(Preprocessor):
|
|
def __init__(self, multiplier):
|
|
self.multiplier = multiplier
|
|
|
|
def transform_batch(self, df):
|
|
return df * self.multiplier
|
|
|
|
|
|
def get_model():
|
|
return tf.keras.Sequential(
|
|
[
|
|
tf.keras.layers.InputLayer(input_shape=()),
|
|
tf.keras.layers.Flatten(),
|
|
tf.keras.layers.Dense(10),
|
|
tf.keras.layers.Dense(1),
|
|
]
|
|
)
|
|
|
|
|
|
def compare_weights(w1: List[ndarray], w2: List[ndarray]) -> bool:
|
|
if not len(w1) == len(w2):
|
|
return False
|
|
size = len(w1)
|
|
for i in range(size):
|
|
comparison = w1[i] == w2[i]
|
|
if not comparison.all():
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
class TestFromModel(unittest.TestCase):
|
|
def setUp(self):
|
|
self.model = get_model()
|
|
self.preprocessor = DummyPreprocessor(1)
|
|
|
|
def test_from_model(self):
|
|
checkpoint = TensorflowCheckpoint.from_model(
|
|
self.model, preprocessor=DummyPreprocessor(1)
|
|
)
|
|
loaded_model = checkpoint.get_model()
|
|
preprocessor = checkpoint.get_preprocessor()
|
|
|
|
assert compare_weights(loaded_model.get_weights(), self.model.get_weights())
|
|
assert preprocessor.multiplier == 1
|
|
|
|
def test_from_saved_model(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model_dir_path = os.path.join(tmp_dir, "my_model")
|
|
self.model.save(model_dir_path, save_format="tf")
|
|
checkpoint = TensorflowCheckpoint.from_saved_model(
|
|
model_dir_path, preprocessor=DummyPreprocessor(1)
|
|
)
|
|
loaded_model = checkpoint.get_model()
|
|
preprocessor = checkpoint.get_preprocessor()
|
|
assert compare_weights(self.model.get_weights(), loaded_model.get_weights())
|
|
assert preprocessor.multiplier == 1
|
|
|
|
def test_from_h5_model(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model_file_path = os.path.join(tmp_dir, "my_model.h5")
|
|
self.model.save(model_file_path)
|
|
checkpoint = TensorflowCheckpoint.from_h5(
|
|
model_file_path, preprocessor=DummyPreprocessor(1)
|
|
)
|
|
loaded_model = checkpoint.get_model()
|
|
preprocessor = checkpoint.get_preprocessor()
|
|
assert compare_weights(self.model.get_weights(), loaded_model.get_weights())
|
|
assert preprocessor.multiplier == 1
|
|
|
|
|
|
def test_tensorflow_checkpoint_saved_model():
|
|
# The test passes if the following can run successfully.
|
|
|
|
def train_fn():
|
|
model = tf.keras.Sequential(
|
|
[
|
|
tf.keras.layers.InputLayer(input_shape=()),
|
|
tf.keras.layers.Flatten(),
|
|
tf.keras.layers.Dense(10),
|
|
tf.keras.layers.Dense(1),
|
|
]
|
|
)
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
model.save(tempdir)
|
|
checkpoint = TensorflowCheckpoint.from_saved_model(tempdir)
|
|
train.report({"my_metric": 1}, checkpoint=checkpoint)
|
|
|
|
trainer = TensorflowTrainer(
|
|
train_loop_per_worker=train_fn, scaling_config=ScalingConfig(num_workers=2)
|
|
)
|
|
|
|
assert trainer.fit().checkpoint
|
|
|
|
|
|
def test_tensorflow_checkpoint_h5():
|
|
# The test passes if the following can run successfully.
|
|
|
|
def train_func():
|
|
model = tf.keras.Sequential(
|
|
[
|
|
tf.keras.layers.InputLayer(input_shape=()),
|
|
tf.keras.layers.Flatten(),
|
|
tf.keras.layers.Dense(10),
|
|
tf.keras.layers.Dense(1),
|
|
]
|
|
)
|
|
with tempfile.TemporaryDirectory() as tempdir:
|
|
model.save(os.path.join(tempdir, "my_model.h5"))
|
|
checkpoint = TensorflowCheckpoint.from_h5(
|
|
os.path.join(tempdir, "my_model.h5")
|
|
)
|
|
train.report({"my_metric": 1}, checkpoint=checkpoint)
|
|
|
|
trainer = TensorflowTrainer(
|
|
train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2)
|
|
)
|
|
|
|
assert trainer.fit().checkpoint
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|