class EmnistDataset(ImageClassificationDataset):
"""
Download and pre-process EMNIST dataset:
python scripts/download.py emnist <desired location>
"""
balance = Param(False)
example_range = Param(None)
class_pool = ''.join([str(i) for i in range(10)] +
[chr(i + ord('A')) for i in range(26)] +
[chr(i + ord('a')) for i in range(26)])
@staticmethod
def sample_classes(n_classes):
classes = np.random.choice(len(EmnistDataset.class_pool),
n_classes,
replace=False)
return [EmnistDataset.class_pool[i] for i in classes]
def _make(self):
param_values = self.param_values()
param_values['one_hot'] = False
param_values['shape'] = param_values['image_shape']
del param_values['image_shape']
x, y, class_map = load_emnist(cfg.data_dir, **param_values)
if x.shape[0] < self.n_examples:
raise Exception("Too few datapoints. Requested {}, "
"only {} are available.".format(
self.n_examples, x.shape[0]))
for _x, _y in x, y:
self._write_example(image=_x, label=class_map[_y])
class ImageClassificationDataset(Dataset):
one_hot = Param()
classes = Param()
image_shape = Param()
include_blank = Param()
_features = None
@property
def features(self):
if self._features is None:
self._features = [
ImageFeature("image", self.obs_shape),
IntegerFeature("label",
self.n_classes if self.one_hot else None),
]
return self._features
@property
def n_classes(self):
return len(self.classes)
@property
def obs_shape(self):
return self.image_shape + (self.depth, )
@property
def action_shape(self):
return self.n_classes if self.one_hot else 1
@property
def depth(self):
return 1
class GridPatchesDataset(PatchesDataset):
grid_shape = Param((2, 2))
spacing = Param((0, 0))
random_offset_range = Param(None)
def _make(self):
self.grid_size = np.product(self.grid_shape)
self.cell_shape = (self.patch_shape[0] + self.spacing[0],
self.patch_shape[1] + self.spacing[1])
return super(GridPatchesDataset, self)._make()
def _sample_patch_locations(self, patch_shapes, **kwargs):
n_patches = len(patch_shapes)
if not n_patches:
return []
indices = np.random.choice(self.grid_size, n_patches, replace=False)
grid_locs = list(zip(*np.unravel_index(indices, self.grid_shape)))
top_left = np.array(grid_locs) * self.cell_shape
if self.random_offset_range is not None:
grid_offset = (
np.random.randint(self.random_offset_range[0]),
np.random.randint(self.random_offset_range[1]),
)
top_left += grid_offset
return [
Rectangle(t, l, m, n)
for (t, l), (m, n, _) in zip(top_left, patch_shapes)
]
class FITSDataset(ImageDataset):
fits_file = Param()
force_memmap = Param()
image_shape = None
"""
Params inherited from ImageDataset, listed here for clarity.
postprocessing = Param("")
tile_shape = Param(None)
n_samples_per_image = Param(1)
n_frames = Param(0)
image_dataset_version = Param(1)
"""
_artifact_names = ['depth']
depth = None
@property
def features(self):
if self._features is None:
annotation_shape = (self.n_frames, -1,
7) if self.n_frames > 0 else (-1, 7)
self._features = [
ImageFeature("image", self.obs_shape, dtype=np.uint16),
VariableShapeArrayFeature("annotations", annotation_shape),
ArrayFeature("offset", (2, ), dtype=np.int32),
StringFeature("filename"),
]
return self._features
def _make(self):
if os.path.isdir(self.fits_file):
local_paths = walk_images(self.fits_file, 'fits')
directory = self.fits_file
else:
local_paths = [os.path.basename(self.fits_file)]
directory = os.path.dirname(self.fits_file)
open_kwargs = dict(memmap=True) if self.force_memmap else {}
for lp in local_paths:
p = os.path.join(directory, lp)
with fits.open(p, **open_kwargs) as hdul:
image = atleast_nd(hdul[0].data, 3)
if self.depth is None:
self.depth = image.shape[2]
self._write_example(
image=image,
annotations=[],
filename=lp,
)
return dict(depth=self.depth)
class NextStep(ConvNet):
kernel_size = Param()
n_channels = Param()
def __init__(self, **kwargs):
layout = [
dict(filters=self.n_channels, kernel_size=self.kernel_size, strides=1, padding="SAME"),
dict(filters=self.n_channels, kernel_size=self.kernel_size, strides=1, padding="SAME"),
]
super(NextStep, self).__init__(layout, check_output_shape=True, **kwargs)
class StochasticGradientDescent(Optimizer):
opt_steps_per_update = Param(1)
sub_batch_size = Param(0)
lr_schedule = Param()
max_grad_norm = Param(None)
noise_schedule = Param(None)
def __init__(self, agents, alg, **kwargs):
super(StochasticGradientDescent, self).__init__(agents)
self.alg = alg
def build_update(self, context):
tvars = self.trainable_variables(for_opt=True)
# `context.objective` is the quantity we want to maximize, but TF minimizes, so use negative.
self.train_op, train_recorded_values = build_gradient_train_op(
-context.objective, tvars, self.alg, self.lr_schedule,
self.max_grad_norm, self.noise_schedule)
context.add_recorded_values(train_recorded_values, train_only=True)
def update(self, n_rollouts, feed_dict, fetches):
sess = tf.get_default_session()
for epoch in range(self.opt_steps_per_update):
record = epoch == self.opt_steps_per_update - 1
if not self.sub_batch_size:
_fetches = [self.train_op, fetches
] if record else self.train_op
fetched = sess.run(_fetches, feed_dict=feed_dict)
else:
for is_final, fd in self.subsample_feed_dict(
n_rollouts, feed_dict):
_fetches = [self.train_op, fetches
] if (record and is_final) else self.train_op
fetched = sess.run(_fetches, feed_dict=fd)
return fetched[1]
def subsample_feed_dict(self, n_rollouts, feed_dict):
updates_per_epoch = int(np.floor(n_rollouts / self.sub_batch_size))
permutation = np.random.permutation(n_rollouts)
offset = 0
for i in range(updates_per_epoch):
indices = permutation[offset:offset + self.sub_batch_size]
fd = {}
for k, v in feed_dict.items():
if isinstance(v, np.ndarray):
fd[k] = v[:, indices, ...]
else:
fd[k] = v
yield i == updates_per_epoch - 1, fd
offset += self.sub_batch_size
class ClassificationNetwork(TensorRecorder):
build_classifier = Param()
loss_type = Param()
def __init__(self, env, updater, scope=None, **kwargs):
self.updater = updater
self.eval_funcs = {}
super(ClassificationNetwork, self).__init__(scope=scope, **kwargs)
def _call(self, inp, label, n_classes, is_training):
self._tensors = dict()
self.maybe_build_subnet('classifier')
logits = self.classifier(inp, n_classes, is_training)
probs = tf.nn.softmax(logits, axis=1)
predicted_labels = tf.cast(tf.argmax(logits, axis=1), label.dtype)
one_hot_label = tf.one_hot(label, n_classes, dtype=probs.dtype)
correct_class_prob = tf.reduce_sum(one_hot_label * probs, axis=1)
if self.loss_type == 'xent':
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=label, logits=logits)
elif self.loss_type.startswith('focal'):
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=label, logits=logits)
gamma = float(self.loss_type.split(',')[1])
focal_weight = (1 - correct_class_prob)**gamma
loss = focal_weight * loss
else:
raise Exception("Unknown loss type: {}".format(self.loss_type))
self.losses = dict(classification=tf.reduce_mean(loss), )
self.record_tensors(
accuracy=tf.equal(predicted_labels, label),
correct_class_prob=correct_class_prob,
)
self._tensors = dict(
logits=logits,
probs=probs,
predicted_labels=predicted_labels,
)
return dict(
tensors=self._tensors,
recorded_tensors=self.recorded_tensors,
losses=self.losses,
)
class ObjectLayer(ScopedFunction):
object_shape = Param()
A = Param()
training_wheels = Param()
noisy = Param()
eval_noisy = Param()
edge_resampler = Param()
obj_temp = Param(help="Higher values -> more uniform")
obj_concrete_temp = Param(help="Higher values -> smoother")
def __init__(self, scope=None, **kwargs):
super().__init__(scope=scope, **kwargs)
self.training_wheels = build_scheduled_value(self.training_wheels, "training_wheels")
self.obj_temp = build_scheduled_value(self.obj_temp, "obj_temp")
self.obj_concrete_temp = build_scheduled_value(self.obj_concrete_temp, "obj_concrete_temp")
def std_nonlinearity(self, std_logit):
# return tf.exp(std)
return (
self._noisy * 2 * tf.nn.sigmoid(tf.clip_by_value(std_logit, -10, 10))
+ (1 - self._noisy) * tf.zeros_like(std_logit)
)
def z_nonlinearity(self, z_logit):
return tf.nn.sigmoid(tf.clip_by_value(z_logit, -10, 10))
@property
def _noisy(self):
return (
self.float_is_training * tf.to_float(self.noisy)
+ (1 - self.float_is_training) * tf.to_float(self.eval_noisy)
)
class SimpleRecurrentRegressionNetwork(ScopedFunction):
use_mask = Param()
cell = None
output_network = None
def _call(self, inp, output_size, is_training):
if self.cell is None:
self.cell = cfg.build_math_cell(scope="regression_cell")
if self.output_network is None:
self.output_network = cfg.build_math_output(scope="math_output")
if self.use_mask:
final_dim = int(inp.shape[-1])
mask, inp = tf.split(inp, (1, final_dim-1), axis=-1)
inp, n_on, _ = apply_mask_and_group_at_front(inp, mask)
else:
batch_size = tf.shape(inp)[0]
n_objects = np.prod(inp.shape[1:-1])
A = inp.shape[-1]
inp = tf.reshape(inp, (batch_size, n_objects, A))
batch_size = tf.shape(inp)[0]
output, final_state = tf.nn.dynamic_rnn(
self.cell, inp, initial_state=self.cell.zero_state(batch_size, tf.float32),
parallel_iterations=1, swap_memory=False, time_major=False)
if self.use_mask:
# Get the output at the end of each sequence.
indices = tf.stack([tf.range(batch_size), n_on-1], axis=1)
output = tf.gather_nd(output, indices)
else:
output = output[:, -1, :]
return self.output_network(output, output_size, is_training)
class ObjKL(Parameterized):
obj_concrete_temp = Param(help="Higher values -> smoother")
def __init__(self, **kwargs):
self.obj_concrete_temp = build_scheduled_value(self.obj_concrete_temp,
"obj_concrete_temp")
super().__init__(**kwargs)
class OmniglotDataset(ImageClassificationDataset):
indices = Param()
@staticmethod
def sample_classes(n_classes):
class_pool = omniglot_classes()
classes = np.random.choice(len(class_pool), n_classes, replace=False)
return [class_pool[i] for i in classes]
def _make(self, **kwargs):
param_values = self.param_values()
param_values['one_hot'] = False
param_values['shape'] = param_values['image_shape']
del param_values['image_shape']
del param_values['n_examples']
x, y, class_map = load_omniglot(cfg.data_dir, **param_values)
if x.shape[0] < self.n_examples:
raise Exception("Too few datapoints. Requested {}, "
"only {} are available.".format(
self.n_examples, x.shape[0]))
for _x, _y in x, y:
self._write_example(image=_x, label=class_map[_y])
class RawDataset(Parameterized):
""" A non-tensorflow dataset, wrapper for data that we might want to cache. """
seed = Param(None)
def __init__(self, **kwargs):
start = time.time()
print("Trying to find dataset in cache...")
directory = kwargs.get(
"data_dir",
os.path.join(cfg.data_dir, "cached_datasets",
self.__class__.__name__))
os.makedirs(directory, exist_ok=True)
params = self.param_values()
param_hash = get_param_hash(params)
print(self.__class__.__name__)
print("Params:")
pprint.pprint(params)
print("Param hash: {}".format(param_hash))
self.directory = os.path.join(directory, str(param_hash))
cfg_filename = os.path.join(self.directory, "config.txt")
if not os.path.exists(cfg_filename):
# Start fresh
try:
shutil.rmtree(self.directory)
except FileNotFoundError:
pass
print("Directory for dataset not found, creating...")
os.makedirs(self.directory, exist_ok=False)
try:
with NumpySeed(self.seed):
self._make()
print("Done creating dataset.")
except BaseException:
try:
shutil.rmtree(self.directory)
except FileNotFoundError:
pass
raise
with open(cfg_filename, 'w') as f:
f.write(pprint.pformat(params))
else:
print("Found.")
print("Took {} seconds.".format(time.time() - start))
print("Features for dataset: ")
pprint.pprint(self.features)
print()
def _make(self):
""" Write data to `self.directory`. """
raise Exception("AbstractMethod.")
class StaticAtariDataset(ReinforcementLearningDataset):
game = Param(aliases="atari_game")
after_warp = Param()
episode_range = Param()
_obs_shape = None
action_dim = 1
reward_dim = 1
rl_data_location = None
@property
def obs_shape(self):
if self._obs_shape is None:
if self.image_shape is not None:
depth = 1 if self.after_warp else 3
self._obs_shape = (*self.image_shape, depth)
else:
if self.postprocessing:
image_shape = self.tile_shape
else:
image_shape = atari_image_shape(self.game, self.after_warp)
if self.after_warp:
self._obs_shape = (*image_shape, 1)
else:
self._obs_shape = (*image_shape, 3)
return self._obs_shape
def _make(self):
directory = os.path.join(cfg.data_dir, "atari_data")
dirs = os.listdir(directory)
game_full_name = "{}NoFrameskip-v4".format(self.game)
starts_with = "atari_data_env={}.datetime=".format(game_full_name)
matching_dirs = [d for d in dirs if d.startswith(starts_with)]
if not matching_dirs:
pprint(sorted(dirs))
raise Exception("No data found for game {}".format(self.game))
directory = os.path.join(directory, sorted(matching_dirs)[-1])
directory = os.path.join(directory,
("after" if self.after_warp else "before") +
"_warp_recording")
scan_recorded_traces(directory, self._callback, self.max_episodes,
self.episode_range)
class ScalorBackground(ScopedFunction):
build_background_encoder = Param()
build_background_decoder = Param()
n_latents_per_channel = Param()
def _call(self, inp, mask, is_training):
self.maybe_build_subnet('background_encoder')
self.maybe_build_subnet('background_decoder')
combined = tf.concat([inp, mask], axis=-1)
latent = self.background_encoder(combined,
2 * self.n_latents_per_channel,
is_training)
mean, std = tf.split(latent, 2, axis=-1)
sample, kl = normal_vae(mean, std, 0, 1)
background = self.background_decoder(sample, None, is_training)
return background, kl