mirror of
https://github.com/vale981/ray
synced 2025-03-10 13:26:39 -04:00
33 lines
781 B
Python
33 lines
781 B
Python
|
# flake8: noqa
|
||
|
|
|
||
|
|
# __use_pretrained_model_start__
|
||
|
|
import ray
|
||
|
|
import tensorflow as tf
|
||
|
|
from ray.air.batch_predictor import BatchPredictor
|
||
|
|
from ray.air.predictors.integrations.tensorflow import (
|
||
|
|
to_air_checkpoint,
|
||
|
|
TensorflowPredictor,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# to simulate having a pretrained model.
|
||
|
|
def build_model() -> tf.keras.Model:
|
||
|
|
model = tf.keras.Sequential(
|
||
|
|
[
|
||
|
|
tf.keras.layers.InputLayer(input_shape=(1,)),
|
||
|
|
tf.keras.layers.Dense(1),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
return model
|
||
|
|
|
||
|
|
|
||
|
|
model = build_model()
|
||
|
|
checkpoint = to_air_checkpoint(model)
|
||
|
|
batch_predictor = BatchPredictor(
|
||
|
|
checkpoint, TensorflowPredictor, model_definition=build_model
|
||
|
|
)
|
||
|
|
predict_dataset = ray.data.range(3)
|
||
|
|
predictions = batch_predictor.predict(predict_dataset)
|
||
|
|
|
||
|
|
# __use_pretrained_model_end__
|