SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 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
1 changed files with 4 additions and 4 deletions
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( freeze_feature_extractor: Optional[bool] = field(
default=True, metadata={"help": "Whether to freeze the feature extractor layers of the model."} 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( verbose_logging: Optional[bool] = field(
default=False, default=False,
metadata={"help": "Whether to log verbose messages or not."}, metadata={"help": "Whether to log verbose messages or not."},
@@ -356,7 +353,6 @@ def main():
config = Wav2Vec2Config.from_pretrained( config = Wav2Vec2Config.from_pretrained(
model_args.model_name_or_path, model_args.model_name_or_path,
cache_dir=model_args.cache_dir, 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": 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)) 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( data_collator = FlaxDataCollatorForWav2Vec2Pretraining(
model=model, feature_extractor=feature_extractor, pad_to_multiple_of=data_args.pad_to_multiple_of model=model, feature_extractor=feature_extractor, pad_to_multiple_of=data_args.pad_to_multiple_of
) )