Reads a dataset from tfds.
(tfds_name: Text,
tfds_data_dir: Text,
tfds_split: Text,
tfds_skip_decoding_feature: Text,
tfds_as_supervised: bool,
input_context: Optional[tf.distribute.InputContext] = None,
seed: Optional[Union[int, tf.Tensor]] = None,
is_training: bool = False,
cache: bool = False,
cycle_length: Optional[int] = None,
block_length: Optional[int] = None)
| 149 | |
| 150 | |
| 151 | def _read_tfds(tfds_name: Text, |
| 152 | tfds_data_dir: Text, |
| 153 | tfds_split: Text, |
| 154 | tfds_skip_decoding_feature: Text, |
| 155 | tfds_as_supervised: bool, |
| 156 | input_context: Optional[tf.distribute.InputContext] = None, |
| 157 | seed: Optional[Union[int, tf.Tensor]] = None, |
| 158 | is_training: bool = False, |
| 159 | cache: bool = False, |
| 160 | cycle_length: Optional[int] = None, |
| 161 | block_length: Optional[int] = None) -> tf.data.Dataset: |
| 162 | """Reads a dataset from tfds.""" |
| 163 | repeat_filenames = is_training and not cache |
| 164 | read_config = tfds.ReadConfig( |
| 165 | interleave_cycle_length=cycle_length, |
| 166 | interleave_block_length=block_length, |
| 167 | input_context=input_context, |
| 168 | shuffle_seed=seed, |
| 169 | repeat_filenames=repeat_filenames, |
| 170 | # Only assert cardinality when we have a finite dataset. |
| 171 | assert_cardinality=not repeat_filenames, |
| 172 | skip_prefetch=True) |
| 173 | |
| 174 | decoders = {} |
| 175 | if tfds_skip_decoding_feature: |
| 176 | for skip_feature in tfds_skip_decoding_feature.split(','): |
| 177 | decoders[skip_feature.strip()] = tfds.decode.SkipDecoding() |
| 178 | |
| 179 | if tfds_name.startswith('mldataset.'): |
| 180 | dataset = tfds.load(name=tfds_name, |
| 181 | split=tfds_split, |
| 182 | as_supervised=tfds_as_supervised, |
| 183 | decoders=decoders if decoders else None, |
| 184 | read_config=read_config) |
| 185 | else: |
| 186 | builder = tfds.builder(tfds_name, data_dir=tfds_data_dir) |
| 187 | if builder.info.splits: |
| 188 | num_shards = len(builder.info.splits[tfds_split].file_instructions) |
| 189 | else: |
| 190 | # The tfds mock path often does not provide splits. |
| 191 | num_shards = 1 |
| 192 | load_kwargs = dict( |
| 193 | name=tfds_name, download=True, split=tfds_split, |
| 194 | shuffle_files=is_training, as_supervised=tfds_as_supervised, |
| 195 | decoders=decoders if decoders else None) |
| 196 | if tfds_data_dir: |
| 197 | load_kwargs.update({'data_dir': tfds_data_dir}) |
| 198 | |
| 199 | if input_context and num_shards < input_context.num_input_pipelines: |
| 200 | # The number of files in the dataset split is smaller than the number of |
| 201 | # input pipelines. We read the entire dataset first and then shard in the |
| 202 | # host memory. |
| 203 | read_config = dataclasses.replace(read_config, input_context=None) |
| 204 | load_kwargs.update({'read_config': read_config}) |
| 205 | dataset = tfds.load(**load_kwargs) |
| 206 | dataset = dataset.shard(input_context.num_input_pipelines, |
| 207 | input_context.input_pipeline_id) |
| 208 | else: |
no test coverage detected