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)
