| ## @package dataset |
| # Module caffe2.python.dataset |
| """ |
| Implementation of an in-memory dataset with structured schema. |
| |
| Use this to store and iterate through datasets with complex schema that |
| fit in memory. |
| |
| Iterating through entries of this dataset is very fast since the dataset |
| is stored as a set of native Caffe2 tensors, thus no type conversion or |
| deserialization is necessary. |
| """ |
| |
| |
| |
| |
| |
| from caffe2.python import core, workspace |
| from caffe2.python.dataio import Reader, Writer |
| from caffe2.python.schema import ( |
| Struct, from_blob_list, from_column_list, InitEmptyRecord) |
| import numpy as np |
| |
| |
| class _DatasetReader(Reader): |
| def __init__(self, dataset, name, batch_size=1, enforce_batch_size=False): |
| """Don't call this directly. Instead, use dataset.reader()""" |
| Reader.__init__(self, dataset.content()) |
| self.dataset = dataset |
| self.name = name or (dataset.name + '_cursor') |
| self.batch_size = batch_size |
| self.enforce_batch_size = enforce_batch_size |
| self.cursor = None |
| |
| def setup_ex(self, init_net, exit_net): |
| if self.cursor is None: |
| self.cursor = init_net.CreateTreeCursor( |
| [], |
| init_net.NextScopedBlob(self.name), |
| fields=self.dataset.fields) |
| |
| def read(self, read_net): |
| assert self.cursor, 'setup not called.' |
| content = self.dataset.content() |
| with core.NameScope(read_net.NextName(self.name)): |
| fields = read_net.ReadNextBatch( |
| [self.cursor] + content.field_blobs(), |
| content.field_names(), |
| batch_size=self.batch_size, |
| enforce_batch_size=self.enforce_batch_size) |
| fields = core.output_to_list(fields) |
| return (read_net.IsEmpty([fields[0]]), fields) |
| |
| def reset(self, net): |
| net.ResetCursor([self.cursor], []) |
| |
| |
| class _DatasetRandomReader(Reader): |
| def __init__(self, dataset, name, indices, batch_size=1, loop_over=False, |
| enforce_batch_size=False): |
| """Don't call this directly. Instead, use dataset.random_reader()""" |
| Reader.__init__(self, dataset.content()) |
| self.dataset = dataset |
| self.cursor = None |
| self.name = name or (dataset.name + '_cursor') |
| self.indices = indices |
| self.batch_size = batch_size |
| self.loop_over = loop_over |
| self.enforce_batch_size = enforce_batch_size |
| |
| def setup_ex(self, init_net, exit_net): |
| if self.cursor is None: |
| self.cursor = init_net.CreateTreeCursor( |
| [], |
| init_net.NextScopedBlob(self.name), |
| fields=self.dataset.fields) |
| |
| def reset(self, net): |
| net.ResetCursor([self.cursor], []) |
| |
| def computeoffset(self, net): |
| self.reset(net) |
| offsets = net.ComputeOffset( |
| [self.cursor] + self.dataset.content().field_blobs(), |
| 'offsets') |
| self.offsets = offsets |
| |
| def sort_and_shuffle(self, net, sort_by_field=None, |
| shuffle_size=1, batch_size=1): |
| # no sorting by default |
| content = self.dataset.content() |
| sort_by_field_idx = -1 |
| if sort_by_field: |
| assert sort_by_field in content.field_names(), ( |
| 'Must be valid field.') |
| sort_by_field_idx = content.field_names().index(sort_by_field) |
| self.reset(net) |
| |
| indices = net.SortAndShuffle( |
| [self.cursor] + content.field_blobs(), |
| 'indices', |
| sort_by_field_idx=sort_by_field_idx, |
| shuffle_size=shuffle_size, |
| batch_size=batch_size) |
| self.indices = indices |
| |
| def read(self, read_net): |
| assert self.cursor, 'setup_ex not called' |
| assert self.indices, 'sort_and_shuffle not called' |
| assert self.offsets, 'computeoffset not called' |
| content = self.dataset.content() |
| with core.NameScope(read_net.NextName(self.name)): |
| fields = read_net.ReadRandomBatch( |
| [self.cursor, self.indices, self.offsets] + ( |
| content.field_blobs()), |
| content.field_names(), |
| batch_size=self.batch_size, |
| enforce_batch_size=self.enforce_batch_size, |
| loop_over=self.loop_over) |
| fields = core.output_to_list(fields) |
| return (read_net.IsEmpty([fields[0]]), fields) |
| |
| |
| class _DatasetWriter(Writer): |
| def __init__(self, content): |
| """Don't call this directly. Use dataset.writer() instead.""" |
| self._content = content |
| self.mutex = None |
| |
| def setup_ex(self, init_net, exit_net): |
| if self.mutex is None: |
| self.mutex = init_net.CreateMutex([]) |
| |
| def write(self, writer_net, fields): |
| """ |
| Add operations to `net` that append the blobs in `fields` to the end |
| of the dataset. An additional operator will also be added that checks |
| the consistency of the data in `fields` against the dataset schema. |
| |
| Args: |
| writer_net: The net that will contain the Append operators. |
| fields: A list of BlobReference to be appeneded to this dataset. |
| """ |
| assert self.mutex is not None, 'setup not called.' |
| field_blobs = self._content.field_blobs() |
| assert len(fields) == len(field_blobs), ( |
| 'Expected %s fields, got %s.' % (len(field_blobs), len(fields))) |
| writer_net.CheckDatasetConsistency( |
| fields, [], fields=self._content.field_names()) |
| writer_net.AtomicAppend( |
| [self.mutex] + field_blobs + list(fields), |
| field_blobs) |
| |
| def commit(self, finish_net): |
| """Commit is a no-op for an in-memory dataset.""" |
| pass |
| |
| |
| def Const(net, value, dtype=None, name=None): |
| """ |
| Create a 'constant' by first creating an external input in the given |
| net, and then feeding the corresponding blob with its provided value |
| in the current workspace. The name is automatically generated in order |
| to avoid clashes with existing blob names. |
| """ |
| assert isinstance(net, core.Net), 'net must be a core.Net instance.' |
| value = np.array(value, dtype=dtype) |
| blob = net.AddExternalInput(net.NextName(prefix=name)) |
| workspace.FeedBlob(str(blob), value) |
| return blob |
| |
| |
| def execution_step_with_progress(name, init_net, substeps, rows_read): |
| # progress reporter |
| report_net = core.Net('report_net') |
| report_net.Print([rows_read], []) |
| return core.execution_step( |
| name, |
| substeps, |
| report_net=report_net, |
| concurrent_substeps=True, |
| report_interval=5) |
| |
| |
| class Dataset: |
| """Represents an in-memory dataset with fixed schema. |
| |
| Use this to store and iterate through datasets with complex schema that |
| fit in memory. |
| |
| Iterating through entries of this dataset is very fast since the dataset |
| is stored as a set of native Caffe2 tensors, thus no type conversion or |
| deserialization is necessary. |
| """ |
| |
| def __init__(self, fields, name=None): |
| """Create an un-initialized dataset with schema provided by `fields`. |
| |
| Before this dataset can be used, it must be initialized, either by |
| `init_empty` or `init_from_dataframe`. |
| |
| Args: |
| fields: either a schema.Struct or a list of field names in a format |
| compatible with the one described in schema.py. |
| name: optional name to prepend to blobs that will store the data. |
| """ |
| assert isinstance(fields, list) or isinstance(fields, Struct), ( |
| 'fields must be either a Struct or a list of raw field names.') |
| if isinstance(fields, list): |
| fields = from_column_list(fields) |
| self.schema = fields |
| self.fields = fields.field_names() |
| self.field_types = fields.field_types() |
| self.name = name or 'dataset' |
| self.field_blobs = fields.field_blobs() if fields.has_blobs() else None |
| |
| def trim(self, net, multiple_of): |
| """ |
| Trims the contents of this dataset so that the number of records is |
| multiple of the given argument. |
| """ |
| net.TrimDataset( |
| self.field_blobs, |
| self.field_blobs, |
| fields=self.fields, |
| multiple_of=multiple_of) |
| |
| def init_empty(self, init_net): |
| """Initialize the blobs for this dataset with empty values. |
| |
| Empty arrays will be immediately fed into the current workspace, |
| and `init_net` will take those blobs as external inputs. |
| """ |
| self.field_blobs = InitEmptyRecord( |
| init_net, self.schema.clone_schema()).field_blobs() |
| |
| def init_from_dataframe(self, net, dataframe): |
| """Initialize the blobs for this dataset from a Pandas dataframe. |
| |
| Each column of the dataframe will be immediately fed into the current |
| workspace, and the `net` will take this blobs as external inputs. |
| """ |
| assert len(self.fields) == len(dataframe.columns) |
| self.field_blobs = [ |
| Const(net, dataframe.as_matrix([col]).flatten(), name=field) |
| for col, field in enumerate(self.fields)] |
| |
| def get_blobs(self): |
| """ |
| Return the list of BlobReference pointing to the blobs that contain |
| the data for this dataset. |
| """ |
| assert self |
| return self.field_blobs |
| |
| def content(self): |
| """ |
| Return a Record of BlobReferences pointing to the full content of |
| this dataset. |
| """ |
| return from_blob_list(self.schema, self.field_blobs) |
| |
| def field_names(self): |
| """Return the list of field names for this dataset.""" |
| return self.fields |
| |
| def field_types(self): |
| """ |
| Return the list of field dtypes for this dataset. |
| |
| If a list of strings, not a schema.Struct, was passed to the |
| constructor, this will return a list of dtype(np.void). |
| """ |
| return self.field_types |
| |
| def reader(self, init_net=None, cursor_name=None, batch_size=1, |
| enforce_batch_size=False): |
| """Create a Reader object that is used to iterate through the dataset. |
| |
| This will append operations to `init_net` that create a TreeCursor, |
| used to iterate through the data. |
| |
| NOTE: Currently, it is not safe to append to a dataset while reading. |
| |
| Args: |
| init_net: net that will be run once to create the cursor. |
| cursor_name: optional name for the blob containing a pointer |
| to the cursor. |
| batch_size: how many samples to read per iteration. |
| |
| Returns: |
| A _DatasetReader that can be used to create operators that will |
| iterate through the dataset. |
| """ |
| assert self.field_blobs, 'Dataset not initialized.' |
| reader = _DatasetReader(self, cursor_name, batch_size, |
| enforce_batch_size) |
| if init_net is not None: |
| reader.setup_ex(init_net, None) |
| return reader |
| |
| def random_reader(self, init_net=None, indices=None, cursor_name=None, |
| batch_size=1, loop_over=False, enforce_batch_size=False): |
| """Create a Reader object that is used to iterate through the dataset. |
| |
| NOTE: The reader order depends on the order in indices. |
| |
| Args: |
| init_net: net that will be run once to create the cursor. |
| indices: blob of reading order |
| cursor_name: optional name for the blob containing a pointer |
| to the cursor. |
| batch_size: how many samples to read per iteration. |
| loop_over: repeat the dataset indefinitely (in the same order) |
| |
| Returns: |
| A DatasetReader that can be used to create operators that will |
| iterate through the dataset according to indices. |
| """ |
| assert self.field_blobs, 'Dataset not initialized.' |
| reader = _DatasetRandomReader( |
| self, cursor_name, indices, batch_size, loop_over, |
| enforce_batch_size) |
| if init_net is not None: |
| reader.setup_ex(init_net, None) |
| return reader |
| |
| def writer(self, init_net=None): |
| """Create a Writer that can be used to append entries into the dataset. |
| |
| NOTE: Currently, it is not safe to append to a dataset |
| while reading from it. |
| NOTE: Currently implementation of writer is not thread safe. |
| TODO: fixme |
| |
| Args: |
| init_net: net that will be run once in order to create the writer. |
| (currently not used) |
| """ |
| assert self.field_blobs, 'Dataset not initialized.' |
| writer = _DatasetWriter(self.content()) |
| if init_net is not None: |
| writer.setup_ex(init_net, None) |
| return writer |
