demo.py 文件源码

def __init__(self, config):

    self._graph = tf.Graph()

    tf_config = None
    if not config.get("gpu", None):
        tf_config = tf.ConfigProto(device_count={"GPU":0})
    else:
        tf_config = tf.ConfigProto(device_count={"GPU":1})
        tf_config.gpu_options.allow_growth = True
        tf_config.gpu_options.per_process_gpu_memory_fraction=config["gpu_memory_fraction"]

    self._sess = tf.Session(config=tf_config, graph=self._graph)

    with self._sess.graph.as_default():
        graph_def = tf.GraphDef()
        with open(config['model'], 'rb') as file:
            graph_def.ParseFromString(file.read())
        tf.import_graph_def(graph_def, name="")

    self._input_x = self._sess.graph.get_operation_by_name('ph_input_x').outputs[0]
    self._pred = self._sess.graph.get_operation_by_name('predictions').outputs[0]
    self._softmax = self._sess.graph.get_operation_by_name('softmax').outputs[0]

    self._http_app = web.Application(
        handlers=[
            web.URLSpec(r"/api/echo/(.*)", EchoHandler, dict(app=self)),
            web.URLSpec(r"/api/image", EchoHandler, dict(app=self)),
            web.URLSpec(r"/ui/segmentation", TestUIHandler, dict(app=self))
        ],
        debug=config["debug"],
    )