| 320 | |
| 321 | |
| 322 | def load_data( |
| 323 | hyperparams: Dict, model_name: str = "resnet_v1" |
| 324 | ) -> [tf.data.Dataset, tf.data.Dataset]: |
| 325 | """ Loads ImageNet data in `tfrecord` format (requires manual data download). |
| 326 | |
| 327 | Args: |
| 328 | hyperparams (Dict): dictionary with necessary hyper-parameters for data loading. |
| 329 | model_name (str): Model name, used to decide which input pre-processing is needed. |
| 330 | Options={supported_model_names}. |
| 331 | |
| 332 | Returns: |
| 333 | train_batches (tf.data.Dataset): 'train' dataset. |
| 334 | val_batches (tf.data.Dataset): 'validation' dataset. |
| 335 | """.format( |
| 336 | supported_model_names=_SUPPORTED_MODEL_NAMES |
| 337 | ) |
| 338 | |
| 339 | data_batches = load_data_tfrecord_tf( |
| 340 | data_dir=hyperparams["tfrecord_data_dir"], |
| 341 | batch_size=hyperparams["batch_size"], |
| 342 | model_name=model_name, |
| 343 | ) |
| 344 | train_batches, val_batches = (data_batches["train"], data_batches["validation"]) |
| 345 | if hyperparams["train_data_size"] is not None: |
| 346 | train_batches = train_batches.take(hyperparams["train_data_size"]) |
| 347 | if hyperparams["val_data_size"] is not None: |
| 348 | val_batches = val_batches.take(hyperparams["val_data_size"]) |
| 349 | |
| 350 | return train_batches, val_batches |