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