Return dataset.
(data_dir, config)
| 128 | |
| 129 | |
| 130 | def get_datasets(data_dir, config): |
| 131 | """Return dataset.""" |
| 132 | if data_dir is None: |
| 133 | raise ValueError("No supplied data directory") |
| 134 | if not os.path.exists(data_dir): |
| 135 | raise ValueError("Data directory {} does not exist".format(data_dir)) |
| 136 | if config.dataset not in ["cifar-10", "cifar-100"]: |
| 137 | raise ValueError("Unknown dataset {}".format(config.dataset)) |
| 138 | |
| 139 | print("Training on {} dataset.".format(config.dataset)) |
| 140 | sys.stdout.flush() |
| 141 | data_dir = os.path.join(data_dir, config.dataset) |
| 142 | if FLAGS.validate: |
| 143 | # 40k Training set |
| 144 | ds_train = cifar_input.get_ds_from_tfrecords( |
| 145 | data_dir=data_dir, |
| 146 | split="train", |
| 147 | data_aug=True, |
| 148 | batch_size=config.batch_size, |
| 149 | epochs=config.epochs, |
| 150 | shuffle=config.shuffle, |
| 151 | data_format=config.data_format, |
| 152 | dtype=config.dtype, |
| 153 | prefetch=config.batch_size) |
| 154 | # 10k Training set |
| 155 | ds_validation = cifar_input.get_ds_from_tfrecords( |
| 156 | data_dir=data_dir, |
| 157 | split="validation", |
| 158 | data_aug=False, |
| 159 | batch_size=config.eval_batch_size, |
| 160 | epochs=1, |
| 161 | shuffle=False, |
| 162 | data_format=config.data_format, |
| 163 | dtype=config.dtype, |
| 164 | prefetch=config.eval_batch_size) |
| 165 | else: |
| 166 | # 50k Training set |
| 167 | ds_train = cifar_input.get_ds_from_tfrecords( |
| 168 | data_dir=data_dir, |
| 169 | split="train_all", |
| 170 | data_aug=True, |
| 171 | batch_size=config.batch_size, |
| 172 | epochs=config.epochs, |
| 173 | shuffle=config.shuffle, |
| 174 | data_format=config.data_format, |
| 175 | dtype=config.dtype, |
| 176 | prefetch=config.batch_size) |
| 177 | ds_validation = None |
| 178 | |
| 179 | # Always compute loss and accuracy on whole test set |
| 180 | ds_train_one_shot = cifar_input.get_ds_from_tfrecords( |
| 181 | data_dir=data_dir, |
| 182 | split="train_all", |
| 183 | data_aug=False, |
| 184 | batch_size=config.eval_batch_size, |
| 185 | epochs=1, |
| 186 | shuffle=False, |
| 187 | data_format=config.data_format, |