Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorboard/util/op_evaluator.py: 44%

1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""TensorBoard helper routine for TF op evaluator. 

16 

17Requires TensorFlow. 

18""" 

19 

20 

21import threading 

22 

23 

24class PersistentOpEvaluator: 

25 """Evaluate a fixed TensorFlow graph repeatedly, safely, efficiently. 

26 

27 Extend this class to create a particular kind of op evaluator, like an 

28 image encoder. In `initialize_graph`, create an appropriate TensorFlow 

29 graph with placeholder inputs. In `run`, evaluate this graph and 

30 return its result. This class will manage a singleton graph and 

31 session to preserve memory usage, and will ensure that this graph and 

32 session do not interfere with other concurrent sessions. 

33 

34 A subclass of this class offers a threadsafe, highly parallel Python 

35 entry point for evaluating a particular TensorFlow graph. 

36 

37 Example usage: 

38 

39 class FluxCapacitanceEvaluator(PersistentOpEvaluator): 

40 \"\"\"Compute the flux capacitance required for a system. 

41 

42 Arguments: 

43 x: Available power input, as a `float`, in jigawatts. 

44 

45 Returns: 

46 A `float`, in nanofarads. 

47 \"\"\" 

48 

49 def initialize_graph(self): 

50 self._placeholder = tf.placeholder(some_dtype) 

51 self._op = some_op(self._placeholder) 

52 

53 def run(self, x): 

54 return self._op.eval(feed_dict: {self._placeholder: x}) 

55 

56 evaluate_flux_capacitance = FluxCapacitanceEvaluator() 

57 

58 for x in xs: 

59 evaluate_flux_capacitance(x) 

60 """ 

61 

62 def __init__(self): 

63 super().__init__() 

64 self._session = None 

65 self._initialization_lock = threading.Lock() 

66 

67 def _lazily_initialize(self): 

68 """Initialize the graph and session, if this has not yet been done.""" 

69 # TODO(nickfelt): remove on-demand imports once dep situation is fixed. 

70 import tensorflow.compat.v1 as tf 

71 

72 with self._initialization_lock: 

73 if self._session: 

74 return 

75 graph = tf.Graph() 

76 with graph.as_default(): 

77 self.initialize_graph() 

78 # Don't reserve GPU because libpng can't run on GPU. 

79 config = tf.ConfigProto(device_count={"GPU": 0}) 

80 self._session = tf.Session(graph=graph, config=config) 

81 

82 def initialize_graph(self): 

83 """Create the TensorFlow graph needed to compute this operation. 

84 

85 This should write ops to the default graph and return `None`. 

86 """ 

87 raise NotImplementedError( 

88 'Subclasses must implement "initialize_graph".' 

89 ) 

90 

91 def run(self, *args, **kwargs): 

92 """Evaluate the ops with the given input. 

93 

94 When this function is called, the default session will have the 

95 graph defined by a previous call to `initialize_graph`. This 

96 function should evaluate any ops necessary to compute the result 

97 of the query for the given *args and **kwargs, likely returning 

98 the result of a call to `some_op.eval(...)`. 

99 """ 

100 raise NotImplementedError('Subclasses must implement "run".') 

101 

102 def __call__(self, *args, **kwargs): 

103 self._lazily_initialize() 

104 with self._session.as_default(): 

105 return self.run(*args, **kwargs)