From 4f299b2446759d17296550e53dd66ea3ab0b3359 Mon Sep 17 00:00:00 2001 From: Nicholas Broad Date: Wed, 7 Sep 2022 07:46:26 -0400 Subject: [PATCH] Accelerator end training (#18910) * add accelerator.end_training() Some trackers need this to end their runs. * fixup and quality * add space * add space again ?!? --- .../run_image_classification_no_trainer.py | 3 +++ examples/pytorch/language-modeling/run_clm_no_trainer.py | 3 +++ examples/pytorch/language-modeling/run_mlm_no_trainer.py | 3 +++ examples/pytorch/multiple-choice/run_swag_no_trainer.py | 3 +++ .../run_semantic_segmentation_no_trainer.py | 3 +++ examples/pytorch/text-classification/run_glue_no_trainer.py | 3 +++ examples/pytorch/token-classification/run_ner_no_trainer.py | 3 +++ examples/pytorch/translation/run_translation_no_trainer.py | 3 +++ 8 files changed, 24 insertions(+) diff --git a/examples/pytorch/image-classification/run_image_classification_no_trainer.py b/examples/pytorch/image-classification/run_image_classification_no_trainer.py index 7037ab6c82..b6e27de0de 100644 --- a/examples/pytorch/image-classification/run_image_classification_no_trainer.py +++ b/examples/pytorch/image-classification/run_image_classification_no_trainer.py @@ -553,6 +553,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/language-modeling/run_clm_no_trainer.py b/examples/pytorch/language-modeling/run_clm_no_trainer.py index dee0fee8a0..c0fcbbd3ce 100755 --- a/examples/pytorch/language-modeling/run_clm_no_trainer.py +++ b/examples/pytorch/language-modeling/run_clm_no_trainer.py @@ -648,6 +648,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/language-modeling/run_mlm_no_trainer.py b/examples/pytorch/language-modeling/run_mlm_no_trainer.py index 9dd519d11e..2a1951f83d 100755 --- a/examples/pytorch/language-modeling/run_mlm_no_trainer.py +++ b/examples/pytorch/language-modeling/run_mlm_no_trainer.py @@ -693,6 +693,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/multiple-choice/run_swag_no_trainer.py b/examples/pytorch/multiple-choice/run_swag_no_trainer.py index aed2ad8aa9..43dee8bfdb 100755 --- a/examples/pytorch/multiple-choice/run_swag_no_trainer.py +++ b/examples/pytorch/multiple-choice/run_swag_no_trainer.py @@ -637,6 +637,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py index dc1dba2f23..cd30a30daa 100644 --- a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py +++ b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py @@ -662,6 +662,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/text-classification/run_glue_no_trainer.py b/examples/pytorch/text-classification/run_glue_no_trainer.py index 3720c9d09b..fadb014831 100644 --- a/examples/pytorch/text-classification/run_glue_no_trainer.py +++ b/examples/pytorch/text-classification/run_glue_no_trainer.py @@ -590,6 +590,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/token-classification/run_ner_no_trainer.py b/examples/pytorch/token-classification/run_ner_no_trainer.py index 72b9ed3af4..4aee8c7eba 100755 --- a/examples/pytorch/token-classification/run_ner_no_trainer.py +++ b/examples/pytorch/token-classification/run_ner_no_trainer.py @@ -746,6 +746,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) diff --git a/examples/pytorch/translation/run_translation_no_trainer.py b/examples/pytorch/translation/run_translation_no_trainer.py index d16c7d5f1c..92bc97e355 100644 --- a/examples/pytorch/translation/run_translation_no_trainer.py +++ b/examples/pytorch/translation/run_translation_no_trainer.py @@ -728,6 +728,9 @@ def main(): output_dir = os.path.join(args.output_dir, output_dir) accelerator.save_state(output_dir) + if args.with_tracking: + accelerator.end_training() + if args.output_dir is not None: accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model)