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"""Ops for tensor_forest.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20from tensorflow.python import ops 21from tensorflow.python.ops import gen_tensor_forest_ops 22from tensorflow.python.ops import resources 23from tensorflow.python.training import saver 24 25 26class TreeVariableSaveable(saver.BaseSaverBuilder.SaveableObject): 27 """Resource that holds a tree.""" 28 29 def __init__(self, type_name, name, container, config, resource_handle_func, 30 create_op_func, is_initialized_op_func, serialize_op_func, 31 deserialize_op_func): 32 33 with ops.name_scope(name, type_name) as name: 34 self._resource_handle = resource_handle_func( 35 container, shared_name=name, name=name) 36 37 self._is_initialized_op = is_initialized_op_func(self._resource_handle) 38 tensor = serialize_op_func(self._resource_handle) 39 self._create_op = create_op_func(self._resource_handle, config) 40 # slice_spec is useful for saving a slice from a variable. 41 # It's not meaningful the tree variable. So we just pass an empty 42 # value. 43 slice_spec = '' 44 specs = [saver.BaseSaverBuilder.SaveSpec(tensor, slice_spec, name)] 45 super(TreeVariableSaveable, self).__init__(self._resource_handle, specs, 46 name) 47 48 ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self) 49 50 resources.register_resource(self._resource_handle, self._create_op, 51 self._is_initialized_op) 52 self._deserialize_op_func = deserialize_op_func 53 54 def restore(self, restored_tensors, unused_restored_shapes): 55 """Restores the associated tree from 'restored_tensors'. 56 57 Args: 58 restored_tensors: the tensors that were loaded from a checkpoint. 59 unused_restored_shapes: the shapes this object should conform to after 60 restore. Not meaningful for trees. 61 62 Returns: 63 The operation that restores the state of the tree variable. 64 """ 65 with ops.control_dependencies([self._create_op]): 66 return self._deserialize_op_func( 67 self._resource_handle, 68 restored_tensors[0], 69 ) 70 71 @property 72 def resource(self): 73 return self._resource_handle 74 75 76def tree_variable(tree_config, name, container=None): 77 return TreeVariableSaveable( 78 'TreeVariable', name, container, tree_config, 79 gen_tensor_forest_ops.tensor_forest_tree_resource_handle_op, 80 gen_tensor_forest_ops.tensor_forest_create_tree_variable, 81 gen_tensor_forest_ops.tensor_forest_tree_is_initialized_op, 82 gen_tensor_forest_ops.tensor_forest_tree_serialize, 83 gen_tensor_forest_ops.tensor_forest_tree_deserialize).resource 84 85 86class ForestVariables(object): 87 """Resource that holds all trees from a forest.""" 88 89 def __init__(self, params, tree_configs=None): 90 91 self._variables = [] 92 93 for i in range(params.n_trees): 94 tree_config = '' 95 if tree_configs is not None: 96 tree_config = tree_configs[i] 97 self._variables.append(tree_variable( 98 tree_config, 99 'tree-%s' % i, 100 )) 101 102 def __getitem__(self, t): 103 return self._variables[t] 104