| |
| |
|
|
|
|
|
|
|
|
|
|
| import os |
|
|
| from caffe2.python import core |
| from caffe2.python.db_file_reader import DBFileReader |
| from caffe2.python.pipeline import pipe |
| from caffe2.python.task import Cluster, TaskGroup |
|
|
|
|
| class CachedReader(DBFileReader): |
|
|
| default_name_suffix = 'cached_reader' |
|
|
| """Reader with persistent in-file cache. |
| |
| Example usage: |
| cached_reader = CachedReader( |
| reader, |
| db_path='/tmp/cache.db', |
| db_type='LevelDB', |
| ) |
| build_cache_step = cached_reader.build_cache_step() |
| with LocalSession() as session: |
| session.run(build_cache_step) |
| |
| Every time new CachedReader is created, it's expected that |
| db_path exists before calling .setup_ex(...) and .read(...). |
| |
| If db_path doesn't exist, it's expected build_cache_step to be called |
| first to build a cache at db_path. |
| |
| build_cache_step will check existence of provided db_path and in case |
| it's missing will initialize it by reading data from original reader. |
| All consequent attempts to read will ignore original reader |
| (i.e. no additional data will be read from it). |
| |
| Args: |
| original_reader: Reader. |
| If provided, it's the original reader used to build the cache file. |
| db_path: str. |
| |
| Optional Args: |
| db_type: str. DB type of file. A db_type is registed by |
| `REGISTER_CAFFE2_DB(<db_type>, <DB Class>)`. |
| Default to 'LevelDB'. |
| name: str or None. Name of CachedReader. |
| Optional name to prepend to blobs that will store the data. |
| Default to '<db_name>_<default_name_suffix>'. |
| batch_size: int. |
| How many examples are read for each time the read_net is run. |
| Defaults to 100. |
| loop_over: bool. |
| If True given, will go through examples in random order endlessly. |
| Defaults to False. |
| """ |
| def __init__( |
| self, |
| original_reader, |
| db_path, |
| db_type='LevelDB', |
| name=None, |
| batch_size=100, |
| loop_over=False, |
| ): |
| assert original_reader is not None, "original_reader can't be None" |
| self.original_reader = original_reader |
|
|
| super(CachedReader, self).__init__( |
| db_path, |
| db_type, |
| name, |
| batch_size, |
| loop_over, |
| ) |
|
|
| def _init_reader_schema(self, *args, **kwargs): |
| """Prepare the reader schema. |
| |
| Since an original reader is given, |
| use it's schema as ground truth. |
| |
| Returns: |
| schema: schema.Struct. Used in Reader.__init__(...). |
| """ |
| return self.original_reader._schema |
|
|
| def build_cache_step(self, overwrite=False): |
| """Build a step for generating cache DB file. |
| |
| If self.db_path exists and not overwritting, build an empty step. |
| Overwise, build a step as follows. |
| Pipe original reader to the _DatasetWriter, |
| so that dataset field blobs are populated. |
| Then save these blobs into a file. |
| |
| Args: |
| overwrite: bool. If true, ignore the existing file |
| and build a new one overwritting the existing one anyway. |
| |
| Returns: |
| build_cache_step: ExecutionStep. |
| The step to be run for building a cache DB file. |
| """ |
| if os.path.exists(self.db_path) and not overwrite: |
| |
| return core.execution_step('build_step', []) |
|
|
| init_net = core.Net('init') |
| self._init_field_blobs_as_empty(init_net) |
| with Cluster(), core.NameScope(self.name), TaskGroup() as copy_tg: |
| pipe(self.original_reader, self.ds.writer(), num_threads=16) |
| copy_step = copy_tg.to_task().get_step() |
| save_net = core.Net('save') |
| self._save_field_blobs_to_db_file(save_net) |
|
|
| return core.execution_step('build_cache', [init_net, copy_step, save_net]) |
|
|
| def _save_field_blobs_to_db_file(self, net): |
| """Save dataset field blobs to a DB file at db_path""" |
| net.Save( |
| self.ds.get_blobs(), |
| [], |
| db=self.db_path, |
| db_type=self.db_type, |
| blob_name_overrides=self.ds.field_names(), |
| absolute_path=True, |
| ) |
|
|
