Antibody_deep_learning / scripts /tf_savedmodel_helper.py
anzhi2710gmailcom's picture
Upload folder using huggingface_hub
fe8e241 verified
Raw
History Blame Contribute Delete
382 Bytes
import numpy as np
import tensorflow as tf
def predict_saved_model(model_dir, input_name, x):
model = tf.saved_model.load(model_dir)
serving = model.signatures["serving_default"]
x = np.asarray(x, dtype=np.float32)
with tf.device("/GPU:0"):
out = serving(**{input_name: tf.constant(x)})
first_key = list(out.keys())[0]
return out[first_key].numpy()