88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
# =============================================================================
|
|
"""Utility classes for testing checkpointing."""
|
|
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops as ops_lib
|
|
from tensorflow.python.ops import gen_lookup_ops
|
|
from tensorflow.python.training import saver as saver_module
|
|
|
|
|
|
class CheckpointedOp:
|
|
"""Op with a custom checkpointing implementation.
|
|
|
|
Defined as part of the test because the MutableHashTable Python code is
|
|
currently in contrib.
|
|
"""
|
|
|
|
# pylint: disable=protected-access
|
|
def __init__(self, name, table_ref=None):
|
|
if table_ref is None:
|
|
self.table_ref = gen_lookup_ops.mutable_hash_table_v2(
|
|
key_dtype=dtypes.string, value_dtype=dtypes.float32, name=name)
|
|
else:
|
|
self.table_ref = table_ref
|
|
self._name = name
|
|
if not context.executing_eagerly():
|
|
self._saveable = CheckpointedOp.CustomSaveable(self, name)
|
|
ops_lib.add_to_collection(ops_lib.GraphKeys.SAVEABLE_OBJECTS,
|
|
self._saveable)
|
|
|
|
@property
|
|
def name(self):
|
|
return self._name
|
|
|
|
@property
|
|
def saveable(self):
|
|
if context.executing_eagerly():
|
|
return CheckpointedOp.CustomSaveable(self, self.name)
|
|
else:
|
|
return self._saveable
|
|
|
|
def insert(self, keys, values):
|
|
return gen_lookup_ops.lookup_table_insert_v2(self.table_ref, keys, values)
|
|
|
|
def lookup(self, keys, default):
|
|
return gen_lookup_ops.lookup_table_find_v2(self.table_ref, keys, default)
|
|
|
|
def keys(self):
|
|
return self._export()[0]
|
|
|
|
def values(self):
|
|
return self._export()[1]
|
|
|
|
def _export(self):
|
|
return gen_lookup_ops.lookup_table_export_v2(self.table_ref, dtypes.string,
|
|
dtypes.float32)
|
|
|
|
class CustomSaveable(saver_module.BaseSaverBuilder.SaveableObject):
|
|
"""A custom saveable for CheckpointedOp."""
|
|
|
|
def __init__(self, table, name):
|
|
tensors = table._export()
|
|
specs = [
|
|
saver_module.BaseSaverBuilder.SaveSpec(tensors[0], "",
|
|
name + "-keys"),
|
|
saver_module.BaseSaverBuilder.SaveSpec(tensors[1], "",
|
|
name + "-values")
|
|
]
|
|
super(CheckpointedOp.CustomSaveable, self).__init__(table, specs, name)
|
|
|
|
def restore(self, restore_tensors, shapes):
|
|
return gen_lookup_ops.lookup_table_import_v2(
|
|
self.op.table_ref, restore_tensors[0], restore_tensors[1])
|
|
# pylint: enable=protected-access
|