riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

[Gradient checkpoining] Update Wav2Vec scripts (#14036)

Browse Source 
Co-authored-by: Stas Bekman <stas@stason.org>
This commit is contained in:
Antonio Carlos Falcão Petri 2021-11-17 14:37:21 -03:00 committed by GitHub
parent c6c075544d
commit 7544efc92e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -48,9 +48,6 @@ class ModelArguments:
freeze_feature_extractor: Optional[bool] = field(
default=True, metadata={"help": "Whether to freeze the feature extractor layers of the model."}
)
gradient_checkpointing: Optional[bool] = field(
default=False, metadata={"help": "Whether to freeze the feature extractor layers of the model."}
)
verbose_logging: Optional[bool] = field(
default=False,
metadata={"help": "Whether to log verbose messages or not."},
@ -356,7 +353,6 @@ def main():
config = Wav2Vec2Config.from_pretrained(
model_args.model_name_or_path,
cache_dir=model_args.cache_dir,
gradient_checkpointing=model_args.gradient_checkpointing,
)
if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
@ -366,6 +362,10 @@ def main():
model = FlaxWav2Vec2ForPreTraining(config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype))
# Activate gradient checkpointing if needed
if training_args.gradient_checkpointing:
model.gradient_checkpointing_enable()
data_collator = FlaxDataCollatorForWav2Vec2Pretraining(
model=model, feature_extractor=feature_extractor, pad_to_multiple_of=data_args.pad_to_multiple_of
)