[TF 2.0 API Docs] Docstring for tf.train.experimental.enable_mixed_precision_graph_rewrite

[tlkh]

In response to #29241

Improved the docstring for tf.train.experimental.enable_mixed_precision_graph_rewrite:

• Added example for using function
• Added colab notebook to demonstrate speed-up without performance penalty
• Added original graphic for loss scaling (source)
• Added more information about graph rewrite operation
• Added performance guide
• Added exception information
• Added more clarification to loss_scale argument

A gist with the rendered docstring is here for ease of review.

Thank you, any feedback or criticism is welcome.

[perfinion]

Great work! Just a few small things to tweak :)

[perfinion]

Can you change the colours in the image as well? On my monitor the green and blue look fairly similar its hard to distinguish.

[tlkh]

I have made all the requested changes!

[perfinion]

@rthadur Can you make sure the links are updated when this is merged?

[tlkh]

Fixed the pylint 80 char line limit errors.

[tlkh]

Pushed the amended commit with the requested changes.

[recrusader]

@tlkh, when enabling "enable_mixed_precision_graph_rewrite", how about the data type of model network when defining the model?
still use FP32? if it is, how to output the FP16 trained model? Thanks!

[tlkh]

Thanks a lot, tlkh. However, if FP16 trained network cannot be saved, it is difficult to use the network outside tensorflow after finishing the training.
In addition, it seems that ops conversion from FP32 to FP16 can happen on new GPU, like Volta. TF cannot do this conversion for old one.

If you wish to use the network outside of TensorFlow for inference in FP16, typically just converting all the datatypes to FP16 will work. Only for training does there have to be special considerations. There are other toolkits to help with inference optimization outside of TensorFlow, such as TensorRT.

Only Volta and Turing GPUs have the FP16 units and Tensor Cores that can benefit from mixed precision. Older GPUs will not see a benefit, hence the feature is not enabled for them.

@recrusader let's move this discussion to a new GitHub issue if you wish to raise any concerns or suggestions. Keep this thread for discussion specifically about this PR.

[recrusader]

I asked this question in official model. However, no one can give me an answer. Thank you very much! I think that I have got the answer.

[tlkh]

@perfinion @martinwicke requesting review once again.

Changes made:

• image and example notebook is now hosted by NVIDIA
• docstring wording has been improved, and updated/merged with recent changes from master
• docstring added to both enable_mixed_precision_graph_rewrite and enable_mixed_precision_graph_rewrite_v1, with the correct distinction made with which kinds of Optimizer are accepted

Thank you!

[martinwicke]

@MarkDaoust will this render correctly? I see there's no indent on the args and returns sections, I'm wondering whether the indent is required.

[tlkh]

@martinwicke thanks for pointing that out. Let me fix that anyway.
EDIT: fixed the indent!

[MarkDaoust]

Hi @tlkh,

It looks like you miss understood @martinwicke's comment.

All the Args/Returns/Raises blocks need to be indented, or they will not be recognized by our linters, and will render incorrectly on tensorflow.org (They will be interpreted as markdown which will just flatten them into a single paragraph).

Thanks.

[tlkh]

Sorry about that! I fixed it.

[tlkh]

Hi tlkh,

It looks like you miss understood martinwicke's comment.

All the Args/Returns/Raises blocks need to be indented, or they will not be recognized by our linters, and will render incorrectly on tensorflow.org (They will be interpreted as markdown which will just flatten them into a single paragraph).

Thanks.

MarkDaoust Hmm I've addressed the changes but I'm not sure why GitHub is still complaining about "1 change requested". Do you need to approve it as well?

Edit: never mind