Docs: add save and load model docs. (#5397)

* docs: add save and load model docs.

* fix: update code style.

* fix: update code style.

* fix: fix typo and modify wording.

* docs: add save model example to notebook.

* fix: fix typo.
This commit is contained in:
Cengguang Zhang 2022-08-16 15:23:30 +08:00 committed by GitHub
parent 611c9aefdf
commit 6af23499e7

View file

@ -118,6 +118,34 @@ est.shutdown()
print(stats) print(stats)
``` ```
### **Step5: Save and Load the Model**
Orca TF2 Estimator supports two formats to save and load the entire model (**TensorFlow SavedModel and Keras H5 Format**). The recommended format is SavedModel, which is the default format when you use `estimator.save()`.
You could also save the model to Keras H5 format by passing `save_format='h5'` or a filename that ends in `.h5` or `.keras` to `estimator.save()`.
**Note that if you run on Apache Hadoop/YARN cluster, you are recommended to save the model to HDFS and load from HDFS as well.**
**1. SavedModel Format**
```python
# save model in SavedModel format
estimator.save("/tmp/cifar10_model")
# load model
estimator.load("/tmp/cifar10_model")
```
**2. HDF5 format**
```python
# save model in H5 format
estimator.save("/tmp/cifar10_model.h5", save_format='h5')
# load model
estimator.load("/tmp/cifar10_model.h5")
```
That's it, the same code can run seamlessly in your local laptop and to distribute K8s or Hadoop cluster. That's it, the same code can run seamlessly in your local laptop and to distribute K8s or Hadoop cluster.
**Note:** You should call `stop_orca_context()` when your program finishes. **Note:** You should call `stop_orca_context()` when your program finishes.