This repository has been archived by the owner on Mar 11, 2021. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 561
/
dual_net_edge_tpu.py
73 lines (63 loc) · 3 KB
/
dual_net_edge_tpu.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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This file contains an implementation of dual_net.py for Google's EdgeTPU.
It can only be used for inference and requires a specially quantized and
compiled model file.
For more information see https://coral.withgoogle.com
"""
import numpy as np
import features as features_lib
import go
from edgetpu.basic.basic_engine import BasicEngine # pylint: disable=import-error
def extract_agz_features(position):
return features_lib.extract_features(position, features_lib.AGZ_FEATURES)
class DualNetworkEdgeTpu():
"""DualNetwork implementation for Google's EdgeTPU."""
def __init__(self, save_file):
self.engine = BasicEngine(save_file)
self.board_size = go.N
self.output_policy_size = self.board_size**2 + 1
input_tensor_shape = self.engine.get_input_tensor_shape()
expected_input_shape = [1, self.board_size, self.board_size, 17]
if not np.array_equal(input_tensor_shape, expected_input_shape):
raise RuntimeError(
'Invalid input tensor shape {}. Expected: {}'.format(
input_tensor_shape, expected_input_shape))
output_tensors_sizes = self.engine.get_all_output_tensors_sizes()
expected_output_tensor_sizes = [self.output_policy_size, 1]
if not np.array_equal(output_tensors_sizes,
expected_output_tensor_sizes):
raise RuntimeError(
'Invalid output tensor sizes {}. Expected: {}'.format(
output_tensors_sizes, expected_output_tensor_sizes))
def run(self, position):
"""Runs inference on a single position."""
probs, values = self.run_many([position])
return probs[0], values[0]
def run_many(self, positions):
"""Runs inference on a list of position."""
processed = map(extract_agz_features, positions)
probabilities = []
values = []
for state in processed:
assert state.shape == (self.board_size, self.board_size,
17), str(state.shape)
result = self.engine.RunInference(state.flatten())
# If needed you can get the raw inference time from the result object.
# inference_time = result[0] # ms
policy_output = result[1][0:self.output_policy_size]
value_output = result[1][-1]
probabilities.append(policy_output)
values.append(value_output)
return probabilities, values