SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix RESOURCE_EXHAUSTED error when dealing with large datasets in Flax example scripts (#18069)

Browse Source 
* Fix RESOURCE_EXHAUSTED error for large datasets on Flax example scripts

* using np.permutation for creating batch_idx

* train_samples_idx -> training_samples_idx

* fix type hints
This commit is contained in:
Duong A. Nguyen
2022-07-11 20:59:08 +07:00
committed by GitHub
parent ac98a88fbc
commit 1e8140caad
5 changed files with 26 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py
View File

@@ -264,7 +264,7 @@ class FlaxDataCollatorForLanguageModeling:
return inputs, labels
def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
def generate_batch_splits(samples_idx: np.ndarray, batch_size: int) -> np.ndarray:
num_samples = len(samples_idx)
samples_to_remove = num_samples % batch_size
@@ -592,7 +592,8 @@ if __name__ == "__main__":
# ======================== Evaluating ==============================
if step % training_args.eval_steps == 0 and step > 0:
eval_samples_idx = jnp.arange(data_args.num_eval_samples)
# Avoid using jax.numpy here in case of TPU training
eval_samples_idx = np.arange(data_args.num_eval_samples)
eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=1)):

8
examples/research_projects/jax-projects/wav2vec2/run_wav2vec2_pretrain_flax.py
View File

@@ -237,7 +237,7 @@ def write_eval_metric(summary_writer, eval_metrics, step):
summary_writer.scalar(f"eval_{metric_name}", value, step)
def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
def generate_batch_splits(samples_idx: np.ndarray, batch_size: int) -> np.ndarray:
num_samples = len(samples_idx)
samples_to_remove = num_samples % batch_size
@@ -541,7 +541,8 @@ def main():
# Generate an epoch by shuffling sampling indices from the train dataset
num_train_samples = len(vectorized_datasets["train"])
train_samples_idx = jax.random.permutation(input_rng, jnp.arange(num_train_samples))
# Avoid using jax.numpy here in case of TPU training
train_samples_idx = np.random.permutation(np.arange(num_train_samples))
train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
# Gather the indexes for creating the batch and do a training step
@@ -574,7 +575,8 @@ def main():
# ======================== Evaluating ==============================
num_eval_samples = len(vectorized_datasets["validation"])
eval_samples_idx = jnp.arange(num_eval_samples)
# Avoid using jax.numpy here in case of TPU training
eval_samples_idx = np.arange(num_eval_samples)
eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
eval_metrics = []