Skip to content
This repository was archived by the owner on Jul 10, 2025. It is now read-only.

Commit b269f13

Browse files
authored
Update 20191106-tf2-tpu-savedmodel.md
Since this API change has been approved and checked in, I now update this doc with the final accepted design and switch status to 'accepted'.
1 parent 0bfae5b commit b269f13

File tree

1 file changed

+57
-3
lines changed

1 file changed

+57
-3
lines changed

rfcs/20191106-tf2-tpu-savedmodel.md

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
# TPU SavedModel Export API for TF2.x
22

3-
Status | Proposed
3+
Status | Accepted
44
:------------ | :-----------------------------------------------------------
55
**RFC #** | [171](https://github.com/tensorflow/community/pull/171)
66
**Author(s)** | Zhuoran Liu ([email protected]), Youlong Cheng ([email protected])
77
**Sponsor** | Jonathan Hseu ([email protected])
8-
**Updated** | 2019-11-06
8+
**Updated** | 2020-11-06
99

1010
## Objective
1111

@@ -82,7 +82,61 @@ whatever reason, e.g. CPU computations can be parallelized, users don’t have
8282
enough TPU resources, etc. In this case, there has to be a way for them to tell
8383
SavedModel that only ‘dense’ is to run on TPU.</center>
8484

85-
## Design Proposal
85+
## Design
86+
87+
The general idea is to allow users to store a function-alias mapping during
88+
model saving, so that they can refer to the function they want to rewrite for
89+
TPU inference when they use downstream graph transformation tools to rewrite
90+
their models for TPU serving.
91+
92+
This alias mapping mechanism is because a single tf.function can generate many
93+
ConcreteFunctions. If a downstream tool wants to refer to all concrete functions
94+
generated by a single tf.function, it can use the `function_aliases` argument to
95+
store a map from the alias name to all concrete function names.
96+
97+
### Major changes
98+
99+
+ For `tf.saved_model.SaveOptions`: A new slot `function_aliases` is added, to
100+
allow users specify alias of functions they potentially wish to be rewritten
101+
by external graph transformation tools (TPU Inference Converter in this
102+
case);
103+
+ For `MetaInfoDef` in `MetaGraphDef` in `SavedModel`: A new field
104+
`functions_aliases` is added, to store names of FunctionDef mapping to their
105+
aliases.
106+
107+
### User facing API
108+
109+
Users can give `FunctionDef`s they potentially want to rewrite for TPU inference
110+
an alias when saving model:
111+
112+
```python
113+
class MyModel:
114+
@tf.function
115+
def func():
116+
...
117+
@tf.function
118+
def serve():
119+
...
120+
func()
121+
122+
model = MyModel()
123+
signatures = {
124+
'serving_default': model.serve.get_concrete_function(),
125+
}
126+
options = tf.saved_model.SaveOptions(function_aliases={
127+
'my_func': model.func,
128+
})
129+
tf.saved_model.save(model, export_dir, signatures, options)
130+
```
131+
132+
And leverage some model conversion tool to convert their CPU model for TPU
133+
serving:
134+
135+
```python
136+
MyModelConversionTool(input_saved_model, output_saved_model, function_alias='my_func')
137+
```
138+
139+
## Alternative Design Considered
86140

87141
### Caveat
88142

0 commit comments

Comments
 (0)