def test_preprocess_and_read(self): max_length = 30 # Write 2 tfrecord shards csv_path = os.path.join(_TEST_DATA_DIR, 'trembl.csv') data.csv_to_tfrecord(csv_path=csv_path, outdir=self._tmpdir, idx=0, total=2) data.csv_to_tfrecord(csv_path=csv_path, outdir=self._tmpdir, idx=1, total=2) # Construct dataset train_files, test_files = data.get_train_test_files(self._tmpdir) train_ds, test_ds = data.load_dataset( train_files=train_files, test_files=test_files, batch_size=1, shuffle_buffer=1, max_train_length=max_length) # Load CSV manually seqs = [] with tf.gfile.GFile(csv_path) as f: for line in f: print(line) seq = line.strip().split(',')[-1] enc = data.protein_domain.encode([seq], pad=False)[0][:max_length] seqs.append(enc) # Confirm we got the same sequences. for ds_x, target in itertools.zip_longest(iter(train_ds), seqs): ds_x = ds_x._numpy()[0] self.assertAllEqual(target, ds_x[:len(target)]) for ds_x, target in itertools.zip_longest(iter(test_ds), seqs): ds_x = ds_x._numpy()[0] self.assertAllEqual(target, ds_x[:len(target)])
def run_experiment( model_dir, data_dir=None, xid=None, batch_size_per_device=128, eval_frequency=500, checkpoint_frequency=10000, save_checkpoints=True, restore_checkpoint=True, num_eval_steps=None, epochs=None, max_train_steps=1000000, # 1 million max_train_length=512, train_summary_frequency=100, max_eval_length=None, model_cls=models.FlaxLM): """Run experiment. Args: model_dir: Directory to save checkpoints and metrics to. data_dir: Directory to load data. xid: Optional experiment id. batch_size_per_device: Batch size per device. eval_frequency: Steps per eval. checkpoint_frequency: How often to checkpoint. If None, only checkpoint once at end of run. save_checkpoints: If True, checkpoints model according to checkpoint_frequency restore_checkpoint: If True, will restore checkpoint from directory. Useful for robustness to preemption. num_eval_steps: Number of eval steps to take on eval dataset. epochs: Number of train epochs. max_train_steps: Stop training after N steps. max_train_length: Crop training sequences to this length. train_summary_frequency: Frequency to write train metrics. max_eval_length: Maximum eval length. Defaults to max_train_length. model_cls: Model class to use. Returns: FlaxLM resulting from running training. """ if xid is not None: model_dir = os.path.join(model_dir, '%s_l%s' % (str(xid), max_train_length)) tf.enable_v2_behavior() if jax.host_id() == 0: summary_writer = tf_summary.create_file_writer(os.path.join( model_dir, 'metrics'), max_queue=1, flush_millis=1000) train_summary_writer = logging_lib.ScalarSummary(step=None, scope='train/', enable_tf=True, verbose=0) eval_summary_writer = logging_lib.ScalarSummary(step=None, scope='eval/', enable_tf=True, verbose=0) batch_size = batch_size_per_device * jax.local_device_count() max_eval_length = max_eval_length or max_train_length train_files, test_files = data.get_train_valid_files(directory=data_dir) train_ds, eval_ds = data.load_dataset(train_files=train_files, test_files=test_files, batch_size=batch_size, max_train_length=max_train_length, max_eval_length=max_eval_length, shuffle_buffer=16384) with contextlib.ExitStack() as stack: # pylint: disable=using-constant-test if jax.host_id() == 0: # Only need metric writer context manager on host 0. stack.enter_context(summary_writer.as_default()) model = model_cls(domain=data.protein_domain, batch_size=batch_size) if restore_checkpoint: try: model.load_checkpoint(model_dir) except ValueError: # No checkpoint to load -> raises ValueError. pass start_step = model.train_step train_ds = train_ds.repeat(epochs) train_iter = iter(train_ds) train_metrics = [] tick = time.time() if jax.host_id() == 0: _write_gin_configs(os.path.join(model_dir, 'config.gin')) num_evals = 0 for step, batch in zip(range(start_step, max_train_steps), train_iter): batch = jax.tree_map(lambda x: x._numpy(), batch) # pylint: disable=protected-access metrics = model.fit_batch(batch) train_metrics.append(metrics) if jax.host_id() == 0 and ( (save_checkpoints and checkpoint_frequency and step % checkpoint_frequency == 0 and step > 0) or step == max_train_steps - 1): model.save_checkpoint(model_dir) if (step + 1) % train_summary_frequency == 0: summary = evaluation.combine_metrics(train_metrics) logging.info('train in step: %d, loss: %.4f', step, summary['loss']) if jax.host_id() == 0: tock = time.time() steps_per_sec = eval_frequency / (tock - tick) tick = tock train_summary_writer('steps per second', steps_per_sec, step) for key, val in summary.items(): if jnp.isnan(val): raise ValueError(f'NaN in {key} at step {step}.') train_summary_writer(key, val, step) # reset metric accumulation for next evaluation cycle. train_metrics = [] if eval_frequency and (step + 1) % eval_frequency == 0: eval_summary = evaluation.evaluate( model=model, eval_ds=eval_ds, num_eval_steps=num_eval_steps) logging.info('eval in step: %d, loss: %.4f', step, eval_summary['loss']) if jax.host_id() == 0: for key, val in eval_summary.items(): eval_summary_writer(key, val, step) tf_summary.flush() summary_writer.flush() if num_evals == 0: # Write out config on first eval. _write_gin_configs( os.path.join(model_dir, 'config_after_eval.gin')) num_evals += 1 if jax.host_id() == 0: tf_summary.flush() summary_writer.close() _write_gin_configs(os.path.join(model_dir, 'config_end.gin')) return model