Skip to content

Commit d324b5f

Browse files
Alfie-Edwardsgeorgepaw
authored andcommitted
Listing Model subclasses where functional and sequential are listed
Summary: Mentioning Model subclasses everywhere where Functional and Sequential models are mentioned in the docs. TF2.5 Only Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, jackh, #documentation, grahamh, georgep Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep Subscribers: jayniep Maniphest Tasks: T58537 Differential Revision: https://phabricator.sourcevertex.net/D63585
1 parent a8c3a30 commit d324b5f

File tree

1 file changed

+96
-32
lines changed

1 file changed

+96
-32
lines changed

tensorflow/compiler/plugin/poplar/docs/keras_tf2.rst

Lines changed: 96 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,21 @@ For example, if we have a model where each step is of batch size 16 and we set
6666
batch of size 64.
6767

6868
Gradient accumulation can be easily enabled for Keras models created inside of
69-
an ``IPUStrategy`` by calling the
70-
:py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_gradient_accumulation_options`
71-
method for Functional Keras models and the
72-
:py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_gradient_accumulation_options`
73-
method for Sequential Keras models. See the respective method documentation
74-
for more details.
69+
an ``IPUStrategy`` by calling the following methods:
70+
71+
.. table::
72+
:width: 100%
73+
74+
+----------------------+----------------------------------------------------------------------------------------------------------+
75+
| ``Functional`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_gradient_accumulation_options` |
76+
+----------------------+----------------------------------------------------------------------------------------------------------+
77+
| ``Sequential`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_gradient_accumulation_options` |
78+
+----------------------+----------------------------------------------------------------------------------------------------------+
79+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_gradient_accumulation_options` |
80+
+----------------------+----------------------------------------------------------------------------------------------------------+
81+
82+
83+
See the respective API documentation for more details.
7584

7685
.. note::
7786

@@ -144,7 +153,7 @@ time and may be useful when debugging your model.
144153
In :numref:`fig-grouped-pipeline`, :numref:`fig-interleaved-pipeline`
145154
and :numref:`fig-sequential-pipeline`, `T` refers to the number of gradient accumulation
146155
steps per replica. See :numref:`pipelining-options` for how to
147-
specify this value.
156+
specify this value.
148157

149158

150159
A detailed explanation of pipelining can be found in the technical note on `Model parallelism with
@@ -332,22 +341,31 @@ to four different pipeline stages as follows:
332341
Pipelining options
333342
__________________
334343

335-
Pipelining options can be set with
336-
:py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_pipelining_options` for a Sequential model,
337-
:py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_pipelining_options` for a Functional model, and
338-
:py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_pipelining_options` for models which subclass `tf.keras.model`.
344+
Pipelining options can be set with the following methods:
345+
346+
.. table::
347+
:width: 100%
348+
349+
+----------------------+-----------------------------------------------------------------------------------------------+
350+
| ``Functional`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_pipelining_options` |
351+
+----------------------+-----------------------------------------------------------------------------------------------+
352+
| ``Sequential`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_pipelining_options` |
353+
+----------------------+-----------------------------------------------------------------------------------------------+
354+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_pipelining_options` |
355+
+----------------------+-----------------------------------------------------------------------------------------------+
356+
357+
358+
See the respective API documentation for more details.
339359

340360
Gradient accumulation is always used when training a pipelined model (unless using the ``Sequential`` schedule). This means
341361
that you must set the option ``gradient_accumulation_steps_per_replica`` using this API when using the ``Grouped`` or
342362
``Interleaved`` schedule. It is optional when using the ``Sequential`` schedule.
343363

344-
345-
The API documentation for :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_pipelining_options` (Sequential model),
346-
:py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_pipelining_options` (Functional model)
347-
and :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_pipelining_options` (Model subclass) explains that the
348-
additional keyword arguments (``pipelining_kwargs``)
349-
will be forwarded to the :py:func:`tensorflow.python.ipu.pipelining_ops.pipeline` operator (which is used internally -
350-
see :numref:`implementation-details`). Refer to the API documentation for :py:func:`~tensorflow.python.ipu.pipelining_ops.pipeline`
364+
The API documentation for ``set_pipelining_options`` explains that the
365+
additional keyword arguments (``pipelining_kwargs``) will be forwarded to the
366+
:py:func:`tensorflow.python.ipu.pipelining_ops.pipeline` operator
367+
(which is used internally - see :numref:`implementation-details`).
368+
Refer to the API documentation for :py:func:`~tensorflow.python.ipu.pipelining_ops.pipeline`
351369
for details about these arguments.
352370

353371
The code sample below illustrates how options can be set with the `set_pipelining_options` API.
@@ -390,9 +408,22 @@ results.
390408

391409
However, IPU TensorFlow also supports *asynchronous callbacks* by providing a
392410
polling mechanism which allows results to be accessed at the earliest possible
393-
instance. Asynchronous callbacks can be enabled by invoking
394-
:py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_asynchronous_callbacks`
395-
with `True` on your ``Sequential`` or ``Functional`` Keras model.
411+
instance. Asynchronous callbacks can be enabled by passing `True` to the
412+
following methods:
413+
414+
.. table::
415+
:width: 100%
416+
417+
+----------------------+---------------------------------------------------------------------------------------------------+
418+
| ``Functional`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_asynchronous_callbacks` |
419+
+----------------------+---------------------------------------------------------------------------------------------------+
420+
| ``Sequential`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_asynchronous_callbacks` |
421+
+----------------------+---------------------------------------------------------------------------------------------------+
422+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_asynchronous_callbacks` |
423+
+----------------------+---------------------------------------------------------------------------------------------------+
424+
425+
426+
See the respective API documentation for more details.
396427

397428
Configuring Infeeds and Outfeed
398429
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -405,13 +436,33 @@ data to and from the IPU devices when using ``fit()``, ``evaluate()`` and
405436
Instances of ``IPUInfeedQueue`` and ``IPUOutfeedQueue`` can be created with
406437
optional arguments which can affect performance of the model.
407438

408-
For configuring the ``IPUInfeedQueue`` use
409-
:py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_infeed_queue_options`
410-
on your ``Sequential`` or ``Functional`` Keras model.
439+
Use the following methods to configure the ``IPUInfeedQueue`` for your Keras model:
440+
441+
.. table::
442+
:width: 100%
443+
444+
+----------------------+-------------------------------------------------------------------------------------------------+
445+
| ``Functional`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_infeed_queue_options` |
446+
+----------------------+-------------------------------------------------------------------------------------------------+
447+
| ``Sequential`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_infeed_queue_options` |
448+
+----------------------+-------------------------------------------------------------------------------------------------+
449+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_infeed_queue_options` |
450+
+----------------------+-------------------------------------------------------------------------------------------------+
451+
452+
453+
Use the following methods to configure the ``IPUOutfeedQueue`` for your Keras model:
454+
455+
.. table::
456+
:width: 100%
457+
458+
+--------------------+--------------------------------------------------------------------------------------------------+
459+
| ``Functional`` | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension.set_outfeed_queue_options` |
460+
+--------------------+--------------------------------------------------------------------------------------------------+
461+
| ``Sequential`` | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_outfeed_queue_options` |
462+
+--------------------+--------------------------------------------------------------------------------------------------+
463+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension.set_outfeed_queue_options` |
464+
+--------------------+--------------------------------------------------------------------------------------------------+
411465

412-
For configuring the ``IPUOutfeedQueue`` use
413-
:py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_outfeed_queue_options`
414-
on your ``Sequential`` or ``Functional`` Keras model.
415466

416467
For example the ``prefetch_depth`` parameter of the ``IPUInfeedQueue`` and the
417468
``buffer_depth`` parameter of the ``IPUOutfeedQueue`` can be configured as
@@ -429,6 +480,12 @@ Saving and loading Keras models
429480
Saving and loading a Keras model must be done within the IPUStrategy scope in
430481
order to save/load IPU-specific information.
431482

483+
When saving and loading ``Model`` subclasses, make sure to save and restore
484+
class members, such as layers, via the config. This can be done by overriding
485+
the ``get_config`` and ``from_config`` methods. Re-creating members from scratch
486+
can cause errors, as the original members may be restored as part of the
487+
IPU-specific internal state.
488+
432489
.. note::
433490
The arguments `pipelining_kwargs` from :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_pipelining_options` and
434491
`gradient_accumulation_optimizer_kwargs` from :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension.set_gradient_accumulation_options`
@@ -446,8 +503,15 @@ Implementation details
446503
When instantiating a standard TensorFlow Keras model inside the scope of
447504
an `IPUStrategy` instance, it is dynamically injected with additional,
448505
IPU-specific, functions.
449-
This is done through the relevant *IPU Keras extension classes*.
450-
For `tensorflow.keras.Sequential`, IPU-specific extensions exist in
451-
:py:class:`~tensorflow.python.ipu.keras.extensions.SequentialExtension` and for
452-
`tensorflow.keras.Model` in
453-
:py:class:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension`.
506+
This is done through the relevant *IPU Keras extension classes*:
507+
508+
.. table::
509+
:width: 100%
510+
511+
+----------------------+------------------------------------------------------------------------+
512+
| ``Functional`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.FunctionalExtension` |
513+
+----------------------+------------------------------------------------------------------------+
514+
| ``Sequential`` model | :py:meth:`~tensorflow.python.ipu.keras.extensions.SequentialExtension` |
515+
+----------------------+------------------------------------------------------------------------+
516+
| ``Model`` subclass | :py:meth:`~tensorflow.python.ipu.keras.extensions.ModelExtension` |
517+
+----------------------+------------------------------------------------------------------------+

0 commit comments

Comments
 (0)