-
Notifications
You must be signed in to change notification settings - Fork 10
/
processes_server.py
56 lines (44 loc) · 1.39 KB
/
processes_server.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import json
from concurrent.futures import ThreadPoolExecutor as Executor
from aiohttp import web
import numpy as np
from benchmarks.compare_tf_optimizations import BinaryModel
from misc.constants import TENSORFLOW_SAVES_DIR
from models.tensorflow_model import Model
N_PARALLEL = 4
def task(model, batch):
pred = model.predict(batch)
str_pred = json.dumps(pred.tolist())
return str_pred
async def index(request):
while True:
try:
model = request.app["models"].pop()
break
except IndexError:
pass
batch = request.app["batch"]
executor = request.app["executor"]
future = executor.submit(task, model, batch)
str_pred = future.result()
request.app["models"].append(model)
return web.Response(text=str_pred)
def build_app():
app = web.Application()
app.router.add_get('/', index)
models = []
for i in range(N_PARALLEL):
models.append(BinaryModel(
saves_dir=TENSORFLOW_SAVES_DIR,
model_file='optimized_graph.pb',
input_node_name=Model.input_node_name,
output_node_name=Model.output_node_name
))
batch = np.random.random((32, 224, 224, 3))
app["models"] = models
app["batch"] = batch
app["executor"] = Executor(N_PARALLEL)
return app
if __name__ == '__main__':
app = build_app()
web.run_app(app, host='0.0.0.0', port=8080)