494 lines
16 KiB
Python
494 lines
16 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.
|
||
|
|
# ==============================================================================
|
||
|
|
|
||
|
|
# pylint: disable=line-too-long
|
||
|
|
"""Inputs and Readers.
|
||
|
|
|
||
|
|
See the @{$python/io_ops} guide.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import absolute_import
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
from tensorflow.python.eager import context
|
||
|
|
from tensorflow.python.framework import dtypes
|
||
|
|
from tensorflow.python.framework import ops
|
||
|
|
from tensorflow.python.lib.io import python_io
|
||
|
|
from tensorflow.python.ops import gen_data_flow_ops
|
||
|
|
from tensorflow.python.ops import gen_io_ops
|
||
|
|
# go/tf-wildcard-import
|
||
|
|
# pylint: disable=wildcard-import
|
||
|
|
from tensorflow.python.ops.gen_io_ops import *
|
||
|
|
from tensorflow.python.util.tf_export import tf_export
|
||
|
|
# pylint: enable=wildcard-import
|
||
|
|
|
||
|
|
|
||
|
|
# pylint: disable=protected-access
|
||
|
|
def _save(filename, tensor_names, tensors, tensor_slices=None, name="save"):
|
||
|
|
"""Save a list of tensors to a file with given names.
|
||
|
|
|
||
|
|
Example usage without slice info:
|
||
|
|
Save("/foo/bar", ["w", "b"], [w, b])
|
||
|
|
|
||
|
|
Example usage with slices:
|
||
|
|
Save("/foo/bar", ["w", "w"], [slice0, slice1],
|
||
|
|
tensor_slices=["4 10 0,2:-", "4 10 2,2:-"])
|
||
|
|
|
||
|
|
Args:
|
||
|
|
filename: the file name of the sstable.
|
||
|
|
tensor_names: a list of strings.
|
||
|
|
tensors: the list of tensors to be saved.
|
||
|
|
tensor_slices: Optional list of strings to specify the shape and slices of
|
||
|
|
a larger virtual tensor that each tensor is a part of. If not specified
|
||
|
|
each tensor is saved as a full slice.
|
||
|
|
name: string. Optional name for the op.
|
||
|
|
|
||
|
|
Requires:
|
||
|
|
The length of tensors should match the size of tensor_names and of
|
||
|
|
tensor_slices.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An Operation that saves the tensors.
|
||
|
|
"""
|
||
|
|
if tensor_slices is None:
|
||
|
|
return gen_io_ops.save(filename, tensor_names, tensors, name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.save_slices(filename, tensor_names, tensor_slices,
|
||
|
|
tensors, name=name)
|
||
|
|
|
||
|
|
|
||
|
|
def _restore_slice(file_pattern, tensor_name, shape_and_slice, tensor_type,
|
||
|
|
name="restore_slice", preferred_shard=-1):
|
||
|
|
"""Restore a tensor slice from a set of files with a given pattern.
|
||
|
|
|
||
|
|
Example usage:
|
||
|
|
RestoreSlice("/foo/bar-?????-of-?????", "w", "10 10 0,2:-", DT_FLOAT)
|
||
|
|
|
||
|
|
Args:
|
||
|
|
file_pattern: the file pattern used to match a set of checkpoint files.
|
||
|
|
tensor_name: the name of the tensor to restore.
|
||
|
|
shape_and_slice: the shape-and-slice spec of the slice.
|
||
|
|
tensor_type: the type of the tensor to restore.
|
||
|
|
name: string. Optional name for the op.
|
||
|
|
preferred_shard: Int. Optional shard to open first in the checkpoint file.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tensor of type "tensor_type".
|
||
|
|
"""
|
||
|
|
base_type = dtypes.as_dtype(tensor_type).base_dtype
|
||
|
|
return gen_io_ops.restore_slice(
|
||
|
|
file_pattern, tensor_name, shape_and_slice, base_type,
|
||
|
|
preferred_shard, name=name)
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("ReaderBase")
|
||
|
|
class ReaderBase(object):
|
||
|
|
"""Base class for different Reader types, that produce a record every step.
|
||
|
|
|
||
|
|
Conceptually, Readers convert string 'work units' into records (key,
|
||
|
|
value pairs). Typically the 'work units' are filenames and the
|
||
|
|
records are extracted from the contents of those files. We want a
|
||
|
|
single record produced per step, but a work unit can correspond to
|
||
|
|
many records.
|
||
|
|
|
||
|
|
Therefore we introduce some decoupling using a queue. The queue
|
||
|
|
contains the work units and the Reader dequeues from the queue when
|
||
|
|
it is asked to produce a record (via Read()) but it has finished the
|
||
|
|
last work unit.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, reader_ref, supports_serialize=False):
|
||
|
|
"""Creates a new ReaderBase.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
reader_ref: The operation that implements the reader.
|
||
|
|
supports_serialize: True if the reader implementation can
|
||
|
|
serialize its state.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
RuntimeError: If eager execution is enabled.
|
||
|
|
"""
|
||
|
|
if context.executing_eagerly():
|
||
|
|
raise RuntimeError(
|
||
|
|
"Readers are not supported when eager execution is enabled. "
|
||
|
|
"Instead, please use tf.data to get data into your model.")
|
||
|
|
|
||
|
|
self._reader_ref = reader_ref
|
||
|
|
self._supports_serialize = supports_serialize
|
||
|
|
|
||
|
|
@property
|
||
|
|
def reader_ref(self):
|
||
|
|
"""Op that implements the reader."""
|
||
|
|
return self._reader_ref
|
||
|
|
|
||
|
|
def read(self, queue, name=None):
|
||
|
|
"""Returns the next record (key, value) pair produced by a reader.
|
||
|
|
|
||
|
|
Will dequeue a work unit from queue if necessary (e.g. when the
|
||
|
|
Reader needs to start reading from a new file since it has
|
||
|
|
finished with the previous file).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
queue: A Queue or a mutable string Tensor representing a handle
|
||
|
|
to a Queue, with string work items.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tuple of Tensors (key, value).
|
||
|
|
key: A string scalar Tensor.
|
||
|
|
value: A string scalar Tensor.
|
||
|
|
"""
|
||
|
|
if isinstance(queue, ops.Tensor):
|
||
|
|
queue_ref = queue
|
||
|
|
else:
|
||
|
|
queue_ref = queue.queue_ref
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_read_v2(self._reader_ref, queue_ref, name=name)
|
||
|
|
else:
|
||
|
|
# For compatibility with pre-resource queues, create a ref(string) tensor
|
||
|
|
# which can be looked up as the same queue by a resource manager.
|
||
|
|
old_queue_op = gen_data_flow_ops.fake_queue(queue_ref)
|
||
|
|
return gen_io_ops.reader_read(self._reader_ref, old_queue_op, name=name)
|
||
|
|
|
||
|
|
def read_up_to(self, queue, num_records, # pylint: disable=invalid-name
|
||
|
|
name=None):
|
||
|
|
"""Returns up to num_records (key, value) pairs produced by a reader.
|
||
|
|
|
||
|
|
Will dequeue a work unit from queue if necessary (e.g., when the
|
||
|
|
Reader needs to start reading from a new file since it has
|
||
|
|
finished with the previous file).
|
||
|
|
It may return less than num_records even before the last batch.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
queue: A Queue or a mutable string Tensor representing a handle
|
||
|
|
to a Queue, with string work items.
|
||
|
|
num_records: Number of records to read.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A tuple of Tensors (keys, values).
|
||
|
|
keys: A 1-D string Tensor.
|
||
|
|
values: A 1-D string Tensor.
|
||
|
|
"""
|
||
|
|
if isinstance(queue, ops.Tensor):
|
||
|
|
queue_ref = queue
|
||
|
|
else:
|
||
|
|
queue_ref = queue.queue_ref
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_read_up_to_v2(self._reader_ref,
|
||
|
|
queue_ref,
|
||
|
|
num_records,
|
||
|
|
name=name)
|
||
|
|
else:
|
||
|
|
# For compatibility with pre-resource queues, create a ref(string) tensor
|
||
|
|
# which can be looked up as the same queue by a resource manager.
|
||
|
|
old_queue_op = gen_data_flow_ops.fake_queue(queue_ref)
|
||
|
|
return gen_io_ops.reader_read_up_to(self._reader_ref,
|
||
|
|
old_queue_op,
|
||
|
|
num_records,
|
||
|
|
name=name)
|
||
|
|
|
||
|
|
def num_records_produced(self, name=None):
|
||
|
|
"""Returns the number of records this reader has produced.
|
||
|
|
|
||
|
|
This is the same as the number of Read executions that have
|
||
|
|
succeeded.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An int64 Tensor.
|
||
|
|
|
||
|
|
"""
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_num_records_produced_v2(self._reader_ref,
|
||
|
|
name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.reader_num_records_produced(self._reader_ref,
|
||
|
|
name=name)
|
||
|
|
|
||
|
|
def num_work_units_completed(self, name=None):
|
||
|
|
"""Returns the number of work units this reader has finished processing.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
An int64 Tensor.
|
||
|
|
"""
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_num_work_units_completed_v2(self._reader_ref,
|
||
|
|
name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.reader_num_work_units_completed(self._reader_ref,
|
||
|
|
name=name)
|
||
|
|
|
||
|
|
def serialize_state(self, name=None):
|
||
|
|
"""Produce a string tensor that encodes the state of a reader.
|
||
|
|
|
||
|
|
Not all Readers support being serialized, so this can produce an
|
||
|
|
Unimplemented error.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A string Tensor.
|
||
|
|
"""
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_serialize_state_v2(self._reader_ref, name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.reader_serialize_state(self._reader_ref, name=name)
|
||
|
|
|
||
|
|
def restore_state(self, state, name=None):
|
||
|
|
"""Restore a reader to a previously saved state.
|
||
|
|
|
||
|
|
Not all Readers support being restored, so this can produce an
|
||
|
|
Unimplemented error.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
state: A string Tensor.
|
||
|
|
Result of a SerializeState of a Reader with matching type.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The created Operation.
|
||
|
|
"""
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_restore_state_v2(
|
||
|
|
self._reader_ref, state, name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.reader_restore_state(self._reader_ref, state, name=name)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def supports_serialize(self):
|
||
|
|
"""Whether the Reader implementation can serialize its state."""
|
||
|
|
return self._supports_serialize
|
||
|
|
|
||
|
|
def reset(self, name=None):
|
||
|
|
"""Restore a reader to its initial clean state.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The created Operation.
|
||
|
|
"""
|
||
|
|
if self._reader_ref.dtype == dtypes.resource:
|
||
|
|
return gen_io_ops.reader_reset_v2(self._reader_ref, name=name)
|
||
|
|
else:
|
||
|
|
return gen_io_ops.reader_reset(self._reader_ref, name=name)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("ReaderRead")
|
||
|
|
ops.NotDifferentiable("ReaderReadUpTo")
|
||
|
|
ops.NotDifferentiable("ReaderNumRecordsProduced")
|
||
|
|
ops.NotDifferentiable("ReaderNumWorkUnitsCompleted")
|
||
|
|
ops.NotDifferentiable("ReaderSerializeState")
|
||
|
|
ops.NotDifferentiable("ReaderRestoreState")
|
||
|
|
ops.NotDifferentiable("ReaderReset")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("WholeFileReader")
|
||
|
|
class WholeFileReader(ReaderBase):
|
||
|
|
"""A Reader that outputs the entire contents of a file as a value.
|
||
|
|
|
||
|
|
To use, enqueue filenames in a Queue. The output of Read will
|
||
|
|
be a filename (key) and the contents of that file (value).
|
||
|
|
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, name=None):
|
||
|
|
"""Create a WholeFileReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
"""
|
||
|
|
rr = gen_io_ops.whole_file_reader_v2(name=name)
|
||
|
|
super(WholeFileReader, self).__init__(rr, supports_serialize=True)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("WholeFileReader")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("TextLineReader")
|
||
|
|
class TextLineReader(ReaderBase):
|
||
|
|
"""A Reader that outputs the lines of a file delimited by newlines.
|
||
|
|
|
||
|
|
Newlines are stripped from the output.
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
# TODO(josh11b): Support serializing and restoring state.
|
||
|
|
|
||
|
|
def __init__(self, skip_header_lines=None, name=None):
|
||
|
|
"""Create a TextLineReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
skip_header_lines: An optional int. Defaults to 0. Number of lines
|
||
|
|
to skip from the beginning of every file.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
"""
|
||
|
|
rr = gen_io_ops.text_line_reader_v2(skip_header_lines=skip_header_lines,
|
||
|
|
name=name)
|
||
|
|
super(TextLineReader, self).__init__(rr)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("TextLineReader")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("FixedLengthRecordReader")
|
||
|
|
class FixedLengthRecordReader(ReaderBase):
|
||
|
|
"""A Reader that outputs fixed-length records from a file.
|
||
|
|
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
# TODO(josh11b): Support serializing and restoring state.
|
||
|
|
|
||
|
|
def __init__(self,
|
||
|
|
record_bytes,
|
||
|
|
header_bytes=None,
|
||
|
|
footer_bytes=None,
|
||
|
|
hop_bytes=None,
|
||
|
|
name=None,
|
||
|
|
encoding=None):
|
||
|
|
"""Create a FixedLengthRecordReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
record_bytes: An int.
|
||
|
|
header_bytes: An optional int. Defaults to 0.
|
||
|
|
footer_bytes: An optional int. Defaults to 0.
|
||
|
|
hop_bytes: An optional int. Defaults to 0.
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
encoding: The type of encoding for the file. Defaults to none.
|
||
|
|
"""
|
||
|
|
rr = gen_io_ops.fixed_length_record_reader_v2(
|
||
|
|
record_bytes=record_bytes,
|
||
|
|
header_bytes=header_bytes,
|
||
|
|
footer_bytes=footer_bytes,
|
||
|
|
hop_bytes=hop_bytes,
|
||
|
|
encoding=encoding,
|
||
|
|
name=name)
|
||
|
|
super(FixedLengthRecordReader, self).__init__(rr)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("FixedLengthRecordReader")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("TFRecordReader")
|
||
|
|
class TFRecordReader(ReaderBase):
|
||
|
|
"""A Reader that outputs the records from a TFRecords file.
|
||
|
|
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
# TODO(josh11b): Support serializing and restoring state.
|
||
|
|
|
||
|
|
def __init__(self, name=None, options=None):
|
||
|
|
"""Create a TFRecordReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
options: A TFRecordOptions object (optional).
|
||
|
|
"""
|
||
|
|
compression_type = python_io.TFRecordOptions.get_compression_type_string(
|
||
|
|
options)
|
||
|
|
|
||
|
|
rr = gen_io_ops.tf_record_reader_v2(
|
||
|
|
name=name, compression_type=compression_type)
|
||
|
|
super(TFRecordReader, self).__init__(rr)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("TFRecordReader")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("LMDBReader")
|
||
|
|
class LMDBReader(ReaderBase):
|
||
|
|
"""A Reader that outputs the records from a LMDB file.
|
||
|
|
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
def __init__(self, name=None, options=None):
|
||
|
|
"""Create a LMDBReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
options: A LMDBRecordOptions object (optional).
|
||
|
|
"""
|
||
|
|
rr = gen_io_ops.lmdb_reader(name=name)
|
||
|
|
super(LMDBReader, self).__init__(rr)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("LMDBReader")
|
||
|
|
|
||
|
|
|
||
|
|
@tf_export("IdentityReader")
|
||
|
|
class IdentityReader(ReaderBase):
|
||
|
|
"""A Reader that outputs the queued work as both the key and value.
|
||
|
|
|
||
|
|
To use, enqueue strings in a Queue. Read will take the front
|
||
|
|
work string and output (work, work).
|
||
|
|
|
||
|
|
See ReaderBase for supported methods.
|
||
|
|
|
||
|
|
@compatibility(eager)
|
||
|
|
Readers are not compatible with eager execution. Instead, please
|
||
|
|
use `tf.data` to get data into your model.
|
||
|
|
@end_compatibility
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, name=None):
|
||
|
|
"""Create a IdentityReader.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: A name for the operation (optional).
|
||
|
|
"""
|
||
|
|
rr = gen_io_ops.identity_reader_v2(name=name)
|
||
|
|
super(IdentityReader, self).__init__(rr, supports_serialize=True)
|
||
|
|
|
||
|
|
|
||
|
|
ops.NotDifferentiable("IdentityReader")
|