[tune] demo exporting trained models in pbt examples (#6533)

python/ray/tune/examples/pbt_convnet_example.py

@@ -18,6 +18,7 @@ import ray
 from ray import tune
 from ray.tune.schedulers import PopulationBasedTraining
 from ray.tune.util import validate_save_restore
+from ray.tune.trial import ExportFormat
 
 # __tutorial_imports_end__
 
@@ -51,6 +52,14 @@ class PytorchTrainble(tune.Trainable):
     def _restore(self, checkpoint_path):
         self.model.load_state_dict(torch.load(checkpoint_path))
 
+    def _export_model(self, export_formats, export_dir):
+        if export_formats == [ExportFormat.MODEL]:
+            path = os.path.join(export_dir, "exported_convnet.pt")
+            torch.save(self.model.state_dict(), path)
+            return {export_formats[0]: path}
+        else:
+            raise ValueError("unexpected formats: " + str(export_formats))
+
     def reset_config(self, new_config):
         for param_group in self.optimizer.param_groups:
             if "lr" in new_config:
@@ -76,7 +85,6 @@ if __name__ == "__main__":
     # check if PytorchTrainble will save/restore correctly before execution
     validate_save_restore(PytorchTrainble)
     validate_save_restore(PytorchTrainble, use_object_store=True)
-    print("Success!")
 
     # __pbt_begin__
     scheduler = PopulationBasedTraining(
@@ -90,18 +98,30 @@ if __name__ == "__main__":
             # allow perturbations within this set of categorical values
             "momentum": [0.8, 0.9, 0.99],
         })
+
     # __pbt_end__
 
     # __tune_begin__
+    class Stopper:
+        def __init__(self):
+            self.should_stop = False
+
+        def stop(self, trial_id, result):
+            max_iter = 5 if args.smoke_test else 100
+            if not self.should_stop and result["mean_accuracy"] > 0.96:
+                self.should_stop = True
+            return self.should_stop or result["training_iteration"] >= max_iter
+
+    stopper = Stopper()
+
     analysis = tune.run(
         PytorchTrainble,
         name="pbt_test",
         scheduler=scheduler,
         reuse_actors=True,
         verbose=1,
-        stop={
-            "training_iteration": 5 if args.smoke_test else 100,
-        },
+        stop=stopper.stop,
+        export_formats=[ExportFormat.MODEL],
         num_samples=4,
         config={
             "lr": tune.uniform(0.001, 1),