def build_dataset(config, is_training=False, fake_data=False):
"""Returns a tf.data.Dataset with <image, label> pairs.
Args:
config: DataConfig instance.
is_training: Whether to build a dataset for training
(with shuffling and image distortions).
fake_data: If true, use randomly generated data.
Returns:
tf.data.Dataset
"""
if fake_data:
logging.info('Generating fake data for config: %s', config)
return image_data_utils.make_fake_data(CIFAR_SHAPE)
logging.info('Building dataset for config:\n%s', attr.asdict(config))
# TODO(yovadia): Split off a validation set from the training set.
if config.corruption_type and config.corruption_static:
return image_data_utils.make_static_dataset(config,
_get_static_cifar_c)
if config.alt_dataset_name:
all_images, all_labels = _download_alt_dataset(config)
else:
all_images, all_labels = _download_cifar_split(config.split,
is_training)
if config.corruption_type:
assert (config.corruption_value
is not None) != (config.corruption_level > 0)
# NOTE: dhtd corruptions expect to be applied before float32 conversion.
apply_corruption = functools.partial(
robustness_dhtd.corrupt,
severity=config.corruption_level,
severity_value=config.corruption_value,
dim=32,
corruption_name=config.corruption_type,
dataset_name='cifar')
all_images = np.stack([apply_corruption(im) for im in all_images])
dataset = tf.data.Dataset.from_tensor_slices((all_images, all_labels))
def prep_fn(image, label):
"""Image preprocessing function."""
if config.roll_pixels:
image = tf.roll(image, config.roll_pixels, -2)
if is_training:
image = tf.image.random_flip_left_right(image)
image = tf.pad(image, [[4, 4], [4, 4], [0, 0]])
image = tf.image.random_crop(image, CIFAR_SHAPE)
image = tf.image.convert_image_dtype(image, tf.float32)
return image, label
return dataset.map(prep_fn)
def build_dataset(config,
batch_size,
is_training=False,
fake_data=False,
use_bfloat16=False):
"""Returns a tf.data.Dataset with <image, label> pairs.
Args:
config: DataConfig instance.
batch_size: Dataset batch size.
is_training: Whether to build a dataset for training
(with shuffling and image distortions).
fake_data: If True, use randomly generated data.
use_bfloat16: If True, use bfloat16. If False, use float32.
Returns:
tf.data.Dataset
"""
if fake_data:
logging.info('Generating fake data for config: %s', config)
return image_data_utils.make_fake_data(IMAGENET_SHAPE).batch(
batch_size)
if config.alt_dataset_name:
dataset = _download_alt_dataset(config, shuffle_files=is_training)
def prep_fn(image_input):
image = tf.image.convert_image_dtype(image_input['image'],
tf.float32)
image = tf.image.crop_to_bounding_box(image, 20, 0, 178, 178)
image = tf.image.resize(image, (224, 224))
# omit CelebA labels
return image, -1
return dataset.map(prep_fn).batch(batch_size)
logging.info('Building dataset for config:\n%s', attr.asdict(config))
if config.corruption_type and config.corruption_static:
return image_data_utils.make_static_dataset(
config, _get_static_imagenet_c).batch(batch_size)
dataset_builder = imagenet_input.ImageNetInput(is_training=is_training,
data_dir=FLAGS.imagenet_dir,
batch_size=batch_size,
dataset_split=config.split,
use_bfloat16=use_bfloat16)
dataset = dataset_builder.input_fn()
if config.corruption_type:
assert (config.corruption_value
is not None) != (config.corruption_level > 0)
# NOTE: dhtd corruptions expect to be applied before float32 conversion.
def apply_corruption(image, label):
"""Apply the corruption function to the image."""
image = tf.image.convert_image_dtype(image, tf.uint8)
corruption_fn = functools.partial(
robustness_dhtd.corrupt,
severity=config.corruption_level,
severity_value=config.corruption_value,
dim=224,
corruption_name=config.corruption_type,
dataset_name='imagenet')
def apply_to_batch(ims):
ims_numpy = ims.numpy()
for i in range(ims_numpy.shape[0]):
ims_numpy[i] = corruption_fn(ims_numpy[i])
return ims_numpy
image = tf.py_function(func=apply_to_batch,
inp=[image],
Tout=tf.float32)
image = tf.clip_by_value(image, 0., 255.) / 255.
return image, label
dataset = dataset.map(apply_corruption)
if config.roll_pixels:
def roll_fn(image, label):
"""Function to roll pixels."""
image = tf.roll(image, config.roll_pixels, -2)
return image, label
dataset = dataset.map(roll_fn)
return dataset