@@ -66,12 +66,21 @@ For example, if we have a model where each step is of batch size 16 and we set
6666batch of size 64.
6767
6868Gradient 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
150159A 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:
332341Pipelining 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
340360Gradient accumulation is always used when training a pipelined model (unless using the ``Sequential `` schedule). This means
341361that 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 `
351369for details about these arguments.
352370
353371The code sample below illustrates how options can be set with the `set_pipelining_options ` API.
@@ -390,9 +408,22 @@ results.
390408
391409However, IPU TensorFlow also supports *asynchronous callbacks * by providing a
392410polling 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
397428Configuring Infeeds and Outfeed
398429~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -405,13 +436,33 @@ data to and from the IPU devices when using ``fit()``, ``evaluate()`` and
405436Instances of ``IPUInfeedQueue `` and ``IPUOutfeedQueue `` can be created with
406437optional 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
416467For 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
429480Saving and loading a Keras model must be done within the IPUStrategy scope in
430481order 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
446503When instantiating a standard TensorFlow Keras model inside the scope of
447504an `IPUStrategy ` instance, it is dynamically injected with additional,
448505IPU-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