Skip to content

Instantly share code, notes, and snippets.

@cutewalker
Last active January 19, 2020 18:03
Show Gist options
  • Save cutewalker/58e1c4f71b5af4822bc732fd619ebda3 to your computer and use it in GitHub Desktop.
Save cutewalker/58e1c4f71b5af4822bc732fd619ebda3 to your computer and use it in GitHub Desktop.
gen tf serving platform config file
# -*- coding: utf-8 -*-
import tensorflow as tf
from tensorflow_serving.config import platform_config_pb2
from tensorflow_serving.servables.tensorflow import session_bundle_config_pb2
from tensorflow_serving.servables.tensorflow import saved_model_bundle_source_adapter_pb2
session_config = tf.ConfigProto()
# config whatever you want
session_config.gpu_options.allow_growth = True
session_config.gpu_options.per_process_gpu_memory_fraction = 0.4
legacy_config=session_bundle_config_pb2.SessionBundleConfig(session_config=session_config)
adapter = saved_model_bundle_source_adapter_pb2.SavedModelBundleSourceAdapterConfig(legacy_config=legacy_config)
config_map = platform_config_pb2.PlatformConfigMap()
config_map.platform_configs['tensorflow'].source_adapter_config.Pack(adapter)
print(config_map)
# run at root of tensorflow_serving repo
TARGET_DIR="$1"
python -m grpc.tools.protoc \
-I . -I ./tensorflow \
--python_out "$TARGET_DIR" \
tensorflow_serving/servables/tensorflow/saved_model_bundle_source_adapter.proto \
tensorflow_serving/servables/tensorflow/session_bundle_config.proto \
tensorflow_serving/config/platform_config.proto
pushd $TARGET_DIR
touch tensorflow_serving/__init__.py
touch tensorflow_serving/config/__init__.py
touch tensorflow_serving/servables/__init__.py
touch tensorflow_serving/servables/tensorflow/__init__.py
popd
platform_configs {
key: "tensorflow"
value {
source_adapter_config {
[type.googleapis.com/tensorflow.serving.SavedModelBundleSourceAdapterConfig] {
legacy_config {
session_config {
gpu_options {
per_process_gpu_memory_fraction: 0.4
allow_growth: true
}
}
}
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment