def main(input: str, noise: float, num_output: int): dset = load_dset(input) for n in range(num_output): num_corrupted_labels = 0 for i in range(dset.size): # Corrupt factor labels with 0.05 probability. one_hot = np.random.choice([False, True], size=5, p=[1 - noise, noise]) for j, x in enumerate(one_hot): if x: fake_factor_value = np.random.uniform(0, 1) dset.factors[i, j] = fake_factor_value num_corrupted_labels += 1 print(f'Corrupted {num_corrupted_labels}/{dset.size} labels.') # Save to file. temp_file = f'/tmp/generate_corrupted_dataset_output-{n}.npz' dset.save(temp_file) print(f'Saved to: {temp_file}') output_prefix = input.split('.npz')[0] + f"-noise{noise}-seed{n}" gfile.copy(temp_file, f'{output_prefix}.npz') gfile.remove(temp_file) print(f'Saved to: {output_prefix}.npz')
def save_checkpoint(tree: Params, path: str, step_for_copy: Optional[int] = None) -> None: """Saves the values of JAX pytrees to disk in a NumPy `.npz` file. Args: tree: A JAX pytree to be saved. path: A path to save the checkpoint. step_for_copy: Optional integer that, when not None, will be used to save a copy of the checkpoint with the name `path-{step_for_copy}`. """ # NOTE: In general, this could be greatly simplified as follows. However, we # currently need to store the leaf names as well in order to be able to load # and reconstruct the tree directly from the checkpoint when initialized a # subset of a model from a pretrained model for fine tuning. # ``` # values, _ = jax.tree_util.tree_flatten(tree) # io_buffer = io.BytesIO() # np.savez(io_buffer, *values) # ``` names_and_vals, _ = _tree_flatten_with_names(tree) io_buffer = io.BytesIO() np.savez(io_buffer, **{k: v for k, v in names_and_vals}) # In order to be robust to interruptions during saving, we first save the # checkpoint to a temporary file, and then rename it to the actual path name. path_tmp = path + "-TEMPORARY" with gfile.GFile(path_tmp, "wb") as f: f.write(io_buffer.getvalue()) gfile.rename(path_tmp, path, overwrite=True) if step_for_copy is not None: gfile.copy(path, f"{path}-{step_for_copy:09d}", overwrite=True)
def load_image(path: str) -> Image: with tempfile.TemporaryDirectory() as t: tmp_image = os.path.join(t, os.path.basename(path)) gfile.copy(path, tmp_image) image = Image.open(tmp_image) return image
def copy(self, source: str, destination: str) -> bool: try: gfile.copy(source, destination) return True except Exception as e: # pylint: disable=broad-except logging.error('Error during copy %s', e) return False
def copy(self, source: str, destination: str) -> bool: if gfile.isdir(destination): # gfile requires a file name for copy destination. return gfile.copy(source, os.path.join(destination, os.path.basename(source)), overwrite=True) return gfile.copy(source, destination, overwrite=True)
def copy_to_experiment_dir(config_file): # copy config file to the experiment directory saved_config_file_path = _config_file_path_to_copy(config_file) # HACK: This is for tensorflow bug workaround. # We can remove following 2 lines once it's been resolved in tensorflow # issue link: https://github.com/tensorflow/tensorflow/issues/28508 if gfile.exists(saved_config_file_path): gfile.remove(saved_config_file_path) gfile.copy(config_file, saved_config_file_path)
def upload_to_gcs(local_path, gcs_path): """Upload local file to Google Cloud Storage. Args: local_path: (string) Local file gcs_path: (string) Google Cloud Storage destination Returns: None """ gfile.copy(local_path, gcs_path)
def _copy_file(src_dir: str, dest_dir: str, file_name: str): src_path = src_dir + '/' + file_name dest_path = dest_dir + '/' + file_name for retries in range(0, 10): try: gfile.copy(src_path, dest_path, overwrite=True) logging.info("copy %s->%s succeeded (retry %d)", src_path, dest_path, retries) return except tf.errors.OpError as ex: logging.error("copy %s->%s (retry %d): %s", src_path, dest_path, retries, ex) time.sleep(1.5**retries)
def save_weights(model, path, overwrite=True): """Customized version of Keras model.save_weights(). - Works with both local and remote storage. - Creates intermediate directories if missing. """ tmp_dir = tempfile.mkdtemp() dirname, basename = os.path.split(path) tmp_path = os.path.join(tmp_dir, basename) model.save_weights(tmp_path) gfile.makedirs(dirname) gfile.copy(tmp_path, path, overwrite=overwrite) gfile.remove(tmp_path)
def _savefile(self, path, img, zoom, grid_height): data = make_grid(img, zoom=zoom, grid_height=grid_height) # Writing takes time, and opening a file for writing erases its contents, # so it's better to write to a temporary file and then copy the results. dirname, basename = osp.dirname(path), osp.basename(path) temp_file = osp.join(dirname, '.' + basename) with gfile.GFile(temp_file, 'wb') as f: img = Image.fromarray(data) img.save(f, format='PNG') gfile.copy(temp_file, path, overwrite=True) gfile.remove(temp_file)
def load_weights(model, path, initialize_weights=True): """Customized version of Keras model.load_weights(). - Works with both local and remote storage. - Automatically initialize weights before loading """ if initialize_weights: # cannot use Keras model.build(), because model input is a dict of tensors fake_data = np.zeros([1, model.grid.size_x, model.grid.size_y], dtype=np.float32) init_state = {k: fake_data for k in model.equation.base_keys} model.call(init_state) tmp_dir = tempfile.mkdtemp() _, basename = os.path.split(path) tmp_path = os.path.join(tmp_dir, basename) gfile.copy(path, tmp_path) model.load_weights(tmp_path) gfile.remove(tmp_path)
def save_config_file(config_file, dest_dir): if not gfile.exists(dest_dir): gfile.makedirs(dest_dir) config_file_dest = os.path.join(dest_dir, 'blueoil_config.yaml') # HACK: This is for tensorflow bug workaround. # We can remove following 2 lines once it's been resolved in tensorflow # issue link: https://github.com/tensorflow/tensorflow/issues/28508 if gfile.exists(config_file_dest): gfile.remove(config_file_dest) return gfile.copy(config_file, config_file_dest)
def log_video(writer, video, tb_key, name, step, work_unit_dir, save_raw=False, scale=False): """Save video frames to tensorboard and a file.""" video_raw = video if scale: video = scale_depth(video) if writer is not None: logging.info('Logging video frames') writer.write_images(step, {f'{tb_key}/{name}': make_image_grid(video)}) filename = f'{tb_key}_{name}_{step:05d}.mp4' local_path = os.path.join('/tmp', filename) logging.info('Writing video to %s', local_path) media.write_video(local_path, video, fps=30) wu_path = os.path.join(work_unit_dir, filename) logging.info('Copying video to %s', wu_path) gfile.copy(local_path, wu_path, overwrite=True) gfile.remove(local_path) if save_raw: # save raw floating point values to scale depth properly raw_filename = f'{tb_key}_{name}_{step:05d}.npy' raw_path = os.path.join(work_unit_dir, raw_filename) logging.info('Saving raw video to %s', raw_path) with gfile.GFile(raw_path, 'wb') as raw_f: onp.save(raw_f, video_raw) logging.info('Done logging video.')
def __init__( self, train_env, eval_env, output_dir, policy_trainer_class, n_real_epochs=10, data_eval_frac=0.125, model_train_batch_size=64, n_model_initial_train_steps=1000, n_model_train_steps_per_epoch=1000, simulated_env_problem_class=( simulated_env_problem.SerializedSequenceSimulatedEnvProblem), simulated_batch_size=16, n_simulated_epochs=1000, trajectory_dump_dir=None, initial_trajectory_dir=None, initial_trajectory_mix_prob=0.5, initial_model=None, init_policy_from_world_model=False, **kwargs): super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs) self._policy_dir = os.path.join(output_dir, 'policy') self._model_dir = os.path.join(output_dir, 'model') # Initialize the policy trainer lazily, so in case of initializing the # policy from world model checkpoint, the trainer will try to load the # checkpoint _after_ it's been created in train_model(). self._policy_trainer_fn = functools.partial( policy_trainer_class, train_env=train_env, eval_env=eval_env, output_dir=self._policy_dir, async_mode=self._async_mode, init_policy_from_world_model_output_dir=( self._model_dir if init_policy_from_world_model else None), ) self._policy_trainer = None self._n_real_epochs = n_real_epochs self._model_train_batch_size = model_train_batch_size self._n_model_initial_train_steps = n_model_initial_train_steps self._n_model_train_steps_per_epoch = n_model_train_steps_per_epoch self._data_eval_frac = data_eval_frac gfile.makedirs(self._model_dir) if initial_model is not None: gfile.copy( initial_model, os.path.join(self._model_dir, 'model.pkl'), overwrite=True, ) self._initial_model = initial_model self._initial_trajectories = None self._sim_env = simulated_env_problem_class( batch_size=None, observation_space=train_env.observation_space, action_space=train_env.action_space, reward_range=train_env.reward_range, discrete_rewards=train_env.discrete_rewards, history_stream=None, # TODO(pkozakowski): Support this. output_dir=self._model_dir, ) self._simulated_batch_size = simulated_batch_size self._n_simulated_epochs = n_simulated_epochs # If trajectory_dump_dir is not provided explicitly, save the trajectories # in output_dir. if trajectory_dump_dir is None: trajectory_dump_dir = os.path.join(output_dir, 'trajectories') self._trajectory_dump_root_dir = trajectory_dump_dir self._initial_trajectory_dir = initial_trajectory_dir self._initial_trajectory_mix_prob = initial_trajectory_mix_prob self._summary_writer = jaxboard.SummaryWriter(self._output_dir) self._simple_epoch = 0 self._policy_epoch = 0 self._model_train_step = 0
def copy(src: str, dst: str, overwrite: bool = False) -> str: if gfile.isdir(dst): dst = os.path.join(dst, os.path.basename(src)) gfile.copy(src, dst, overwrite=overwrite) return dst
def main(unused_argv): """Induce and write set of rules.""" examples = tsv_utils.read_tsv(FLAGS.input) config = json_utils.json_file_to_dict(FLAGS.config) if not config.get("allow_duplicate_examples", True): examples = set([tuple(ex) for ex in examples]) examples = sorted(examples, key=lambda e: (len(e[0]), e)) seed_rules = set() # Add mannual seed rules. if FLAGS.seed_rules_file: for seed_rules_file in FLAGS.seed_rules_file: seed_rules |= set(qcfg_file.read_rules(seed_rules_file)) target_grammar_rules = (target_grammar.load_rules_from_file( FLAGS.target_grammar) if FLAGS.target_grammar else None) num_partitions = config.get("num_partitions", 1) partition_to_examples = np.array_split(examples, num_partitions) induction_state = induction_utils.InductionState(FLAGS.output, config) if FLAGS.restore_partition or FLAGS.restore_step: # Restore from an existing induction state. induction_state.restore_state(FLAGS.restore_partition, FLAGS.restore_step) else: # Initialize the induction state with manual seed rules. induction_state.current_rules = seed_rules.copy() while induction_state.current_partition < num_partitions: current_examples = induction_utils.get_examples_up_to_partition( partition_to_examples, induction_state.current_partition) logging.info("Partition: %s, number of examples: %s.", induction_state.current_partition, len(current_examples)) # Restore induction state from the last partition, since the induced # rules are saved with Beam I/O. # TODO(linluqiu): Investigate performance improvements from using a single # Beam pipeline for all partitions without requiring rules to be explicitly # written and read from disk. if (induction_state.current_partition > 0 and induction_state.current_partition != FLAGS.restore_partition): induction_state.restore_state( induction_state.current_partition - 1, config["max_num_steps"] - 1) # At the first step of each partition, we add a rule corresponding to # each example in the partition. if induction_state.current_step_in_partition == 0: induction_state.current_rules |= induction_utils.get_example_rules( partition_to_examples[induction_state.current_partition]) policy = greedy_policy.GreedyPolicy(config, current_examples, seed_rules, target_grammar_rules, verbose=FLAGS.verbose) # pylint: disable=cell-var-from-loop def _induce_rules(pipeline): _ = (pipeline | "InduceRulesForPartition" >> InduceRulesForPartition(policy, induction_state)) # pylint: enable=cell-var-from-loop pipeline_options = beam.options.pipeline_options.PipelineOptions( FLAGS.pipeline_options) with beam.Pipeline(pipeline_options) as pipeline: _induce_rules(pipeline) counters = pipeline.result.metrics().query()["counters"] for counter in counters: logging.info("%s %s: %s", counter.key.step, counter.key.metric.name, counter.committed) rule_file = induction_state.get_rule_file( induction_state.current_partition, config["max_num_steps"] - 1) gfile.copy(rule_file, FLAGS.output, overwrite=True) induction_state.current_partition += 1
def copy(self, source: str, destination: str) -> bool: return gfile.copy(source, destination)
def __init__( self, train_env, eval_env, output_dir, policy_trainer_class, n_real_epochs=10, data_eval_frac=0.125, model_train_batch_size=64, n_model_initial_train_steps=1000, n_model_train_steps_per_epoch=1000, simulated_env_problem_class=( simulated_env_problem.SerializedSequenceSimulatedEnvProblem), simulated_batch_size=16, n_simulated_epochs=1000, trajectory_dump_dir=None, initial_trajectory_dir=None, initial_trajectory_mix_prob=0.5, initial_model=None, **kwargs): super(SimPLe, self).__init__(train_env, eval_env, output_dir, **kwargs) self._policy_dir = os.path.join(output_dir, "policy") self._policy_trainer = policy_trainer_class( train_env=train_env, eval_env=eval_env, output_dir=self._policy_dir, async_mode=self._async_mode, ) self._n_real_epochs = n_real_epochs self._model_train_batch_size = model_train_batch_size self._n_model_initial_train_steps = n_model_initial_train_steps self._n_model_train_steps_per_epoch = n_model_train_steps_per_epoch self._data_eval_frac = data_eval_frac self._model_dir = os.path.join(output_dir, "model") gfile.makedirs(self._model_dir) if initial_model is not None: gfile.copy( initial_model, os.path.join(self._model_dir, "model.pkl"), overwrite=True, ) self._initial_model = initial_model self._sim_env = simulated_env_problem_class( batch_size=None, observation_space=train_env.observation_space, action_space=train_env.action_space, reward_range=train_env.reward_range, discrete_rewards=train_env.discrete_rewards, history_stream=None, # TODO(pkozakowski): Support this. output_dir=self._model_dir, ) self._simulated_batch_size = simulated_batch_size self._n_simulated_epochs = n_simulated_epochs # If trajectory_dump_dir is not provided explicitly, save the trajectories # in output_dir. if trajectory_dump_dir is None: trajectory_dump_dir = os.path.join(output_dir, "trajectories") self._trajectory_dump_root_dir = trajectory_dump_dir self._initial_trajectory_dir = initial_trajectory_dir self._initial_trajectory_mix_prob = initial_trajectory_mix_prob self._summary_writer = jaxboard.SummaryWriter(self._output_dir) self._simple_epoch = 0 self._policy_epoch = 0 self._model_train_step = 0
def copy_remote(src, dst, overwrite=False): """ Allows us to copy a file from local to remote or vice versa """ return gfile.copy(src, dst, overwrite)
def copy_to_experiment_dir(config_file): # copy config file to the experiment directory saved_config_file_path = _config_file_path_to_copy(config_file) gfile.copy(config_file, saved_config_file_path, overwrite=True)