25. Keras Python API

25.1. IPU specific Keras extensions

class keras.ipu.extensions.FunctionalExtension(*args, **kwargs)
get_pipeline_stage_assignment()

Returns the pipeline stage assignment of all the layers in the model.

If set_pipeline_stage_assignment() has been called before, then it returns a copy of the current assignment, otherwise returns a list of FunctionalLayerPipelineStageAssignment for each invocation of each layer in the model (excluding input layers) in post order (which means that layers are returned in the order they are executed).

print_pipeline_stage_assignment_summary(line_length=None, print_fn=None)

Prints a summary of the pipeline stage assignment of the model.

Parameters

  • line_length – Total length of printed lines (e.g. set this to adapt the display to different terminal window sizes).

  • print_fn – Print function to use. It will be called on each line of the summary. You can set it to a custom function in order to capture the string summary. It defaults to print (prints to stdout).

reset_pipeline_stage_assignment()

Resets the pipeline stage assignment so that the model is no longer pipelined.

set_asynchronous_callbacks(asynchronous=False)

Sets the asynchronous callbacks options when calling fit(), evaluate() and predict().

When running fit(), evaluate() and predict() the callbacks the model is configured with are executed after steps_per_execution have executed. Enabling asynchronous callbacks means that the callbacks are invoked after every step, even when steps_per_execution > 1. This can reduce the latency of receiving per step results and metrics at a cost of an extra thread running in the background of the application. Note that this option is ignored for the fit() and evaluate() when running a pipelined model and accumulate_outfeed=True (configured via set_pipelining_options).

Parameters

asynchronous – Whether asynchronous callbacks should be enabled.

set_gradient_accumulation_options(gradient_accumulation_steps_per_replica=None, gradient_accumulation_reduction_method='sum', **gradient_accumulation_optimizer_kwargs)

Sets the gradient accumulation options for non-pipelined models which are to be used when training a model.

When set, and gradient_accumulation_steps_per_replica > 1, the optimizer which the current model has been compiled with is wrapped in GradientAccumulationOptimizerV2. This means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when gradient accumulation is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • reduction_method – Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • gradient_accumulation_optimizer_kwargs – All remaining keyword arguments are forwarded to GradientAccumulationOptimizerV2. See the optimizer for all the available arguments. Must not contain opt or num_mini_batches as keys. Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_gradient_accumulation_options again.

set_infeed_queue_options(**kwargs)

Sets the options for all instances of IPUInfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUInfeedQueue is created to efficiently feed data from the dataset to the device. Instances of IPUInfeedQueue can be created with optional arguments, such as prefetch_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUInfeedQueue.

set_outfeed_queue_options(**kwargs)

Sets the options for all instances of IPUOutfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUOutfeedQueue is created to efficiently feed data from the device to the host. Instances of IPUOutfeedQueue can be created with optional arguments, such as buffer_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUOutfeedQueue.

set_pipeline_stage_assignment(pipeline_stage_assignment)

Sets the pipeline stage assignment of all the invocations of all the layers in the model.

Sets the pipeline stage assignment all the invocations of all the layers (excluding input layers) in the model which is used to create a model-parallel execution of this model when calling fit(), evaluate() and predict(). Note that this pipelining stage assignment is ignored when using the call() function on this model.

Parameters

pipeline_stage_assignment – A list of the same length as the total number of invocations of all the layers in this model (excluding input layers). All elements have to be instances of FunctionalLayerPipelineStageAssignment which are used to indicate which pipeline stage a particular layer invocation should be assigned to.

Raises

ValueError – pipeline_stage_assignment is not a valid assignment.

set_pipelining_options(gradient_accumulation_steps_per_replica=None, device_mapping=None, accumulate_outfeed=None, gradient_accumulation_reduction_method='sum', **pipelining_kwargs)

Sets the pipelining options, including gradient accumulation options, for pipelined models.

Before training a pipelined model, gradient_accumulation_steps_per_replica argument needs to be set as pipelined models always perform gradient accumulation when training. Setting gradient_accumulation_steps_per_replica > 1 means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

When training a data-parallel model, enabling gradient accumulation also reduces the communication overhead as the all-reduce of gradients is now performed after each replica has performed gradient_accumulation_steps_per_replica steps instead of after each step.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when pipelining is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • device_mapping – If provided, a list of length equal to the number of pipeline stages assigned in this model. An element at index i in the list represents which IPU the i’th pipeline stage should reside on. This can be used to make sure computational stages which share Keras layers/tf.Variable objects are resident on the same IPU. This value is saved/loaded when the model is saved/loaded.

  • accumulate_outfeed – The metrics from the model are normally enqueued as soon as they’re available. If this option is True, the data will instead be accumulated when they’re available and enqueued at the end of pipeline execution, reducing the amount of host <-> device communication. When used with training, the accumulated metrics are normalised by gradient_accumulation_steps_per_replica. When used with evaluation, the accumulated metrics are normalised by steps_per_epoch. This option is ignored when doing prediction. When using accumulate_outfeed, model callbacks will be called with the same data for the batches which the data was accumulated for. This value is saved/loaded when the model is saved/loaded.

  • gradient_accumulation_reduction_method – (Experimental) Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • pipelining_kwargs – All remaining keyword arguments are forwarded to pipeline(). Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_pipelining_options again.

class keras.ipu.extensions.FunctionalLayerPipelineStageAssignment(layer, node_index, pipeline_stage=None)

A class to indicate at which pipeline stage a layer in a Functional model should be executed.

Keras Layers can be called multiple times in order to share weights between layers. Each of these calls produces Tensor output which can be executed in different pipeline stages (as long as these stages are mapped to the same device).

property inbound_layers

Returns the input layers for the layer in this assignment. This can be useful for identifying which specific node_index this is.

property layer

Returns the Keras layer associated with this assignment.

property node_index

Returns the specific call to the layer that produced a Tensor.

property pipeline_stage

Returns the pipeline stage this layer has been assigned to. If None, then this layer has not been assigned to a pipeline stage.

class keras.ipu.extensions.ModelExtension(*args, **kwargs)
get_pipeline_stage_assignment()

Returns the pipeline stage assignment of the layers in the model.

If set_pipeline_stage_assignment() has been called before, then it returns a copy of the current assignment, otherwise returns a list of ModelLayerPipelineStageAssignment for each layer in the model in post order (which means that layers are returned in the order they are executed).

print_pipeline_stage_assignment_summary(line_length=None, print_fn=None)

Prints a summary of the pipeline stage assignment of the model.

Parameters

  • line_length – Total length of printed lines (e.g. set this to adapt the display to different terminal window sizes).

  • print_fn – Print function to use. It will be called on each line of the summary. You can set it to a custom function in order to capture the string summary. It defaults to print (prints to stdout).

reset_pipeline_stage_assignment()

Resets the pipeline stage assignment so that the model is no longer pipelined.

set_asynchronous_callbacks(asynchronous=False)

Sets the asynchronous callbacks options when calling fit(), evaluate() and predict().

When running fit(), evaluate() and predict() the callbacks the model is configured with are executed after steps_per_execution have executed. Enabling asynchronous callbacks means that the callbacks are invoked after every step, even when steps_per_execution > 1. This can reduce the latency of receiving per step results and metrics at a cost of an extra thread running in the background of the application. Note that this option is ignored for the fit() and evaluate() when running a pipelined model and accumulate_outfeed=True (configured via set_pipelining_options).

Parameters

asynchronous – Whether asynchronous callbacks should be enabled.

set_gradient_accumulation_options(gradient_accumulation_steps_per_replica=None, gradient_accumulation_reduction_method='sum', **gradient_accumulation_optimizer_kwargs)

Sets the gradient accumulation options for non-pipelined models which are to be used when training a model.

When set, and gradient_accumulation_steps_per_replica > 1, the optimizer which the current model has been compiled with is wrapped in GradientAccumulationOptimizerV2. This means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when gradient accumulation is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • reduction_method – Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • gradient_accumulation_optimizer_kwargs – All remaining keyword arguments are forwarded to GradientAccumulationOptimizerV2. See the optimizer for all the available arguments. Must not contain opt or num_mini_batches as keys. Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_gradient_accumulation_options again.

set_infeed_queue_options(**kwargs)

Sets the options for all instances of IPUInfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUInfeedQueue is created to efficiently feed data from the dataset to the device. Instances of IPUInfeedQueue can be created with optional arguments, such as prefetch_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUInfeedQueue.

set_outfeed_queue_options(**kwargs)

Sets the options for all instances of IPUOutfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUOutfeedQueue is created to efficiently feed data from the device to the host. Instances of IPUOutfeedQueue can be created with optional arguments, such as buffer_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUOutfeedQueue.

set_pipeline_stage_assignment(pipeline_stage_assignment)

Sets the pipeline stage assignment for all the layers in the model.

Sets the pipeline stage assignment of all the layers in the model which is used to create a model-parallel execution of this Model when calling fit(), evaluate() and predict(). Note that this pipelining stage assignment is ignored when using the call() function on this model.

Parameters

pipeline_stage_assignment – A list of the same length as the number of layers in this model. All elements can be either intergers or instances of ModelLayerPipelineStageAssignment. If all the elements are integers, then a layer in this model at index i is assigned to a pipeline stage pipeline_stage_assignment[i]. Otherwise, if all the elements are of type ModelLayerPipelineStageAssignment then a layer in this model at index i is assigned to a pipeline stage indicated by pipeline_stage_assignment[i].pipeline_stage.

Raises

ValueError – pipeline_stage_assignment is not a valid assignment.

set_pipelining_options(gradient_accumulation_steps_per_replica=None, device_mapping=None, accumulate_outfeed=None, gradient_accumulation_reduction_method='sum', **pipelining_kwargs)

Sets the pipelining options, including gradient accumulation options, for pipelined models.

Before training a pipelined model, gradient_accumulation_steps_per_replica argument needs to be set as pipelined models always perform gradient accumulation when training. Setting gradient_accumulation_steps_per_replica > 1 means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

When training a data-parallel model, enabling gradient accumulation also reduces the communication overhead as the all-reduce of gradients is now performed after each replica has performed gradient_accumulation_steps_per_replica steps instead of after each step.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when pipelining is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • device_mapping – If provided, a list of length equal to the number of pipeline stages assigned in this model. An element at index i in the list represents which IPU the i’th pipeline stage should reside on. This can be used to make sure computational stages which share Keras layers/tf.Variable objects are resident on the same IPU. This value is saved/loaded when the model is saved/loaded.

  • accumulate_outfeed – The metrics from the model are normally enqueued as soon as they’re available. If this option is True, the data will instead be accumulated when they’re available and enqueued at the end of pipeline execution, reducing the amount of host <-> device communication. When used with training, the accumulated metrics are normalised by gradient_accumulation_steps_per_replica. When used with evaluation, the accumulated metrics are normalised by steps_per_epoch. This option is ignored when doing prediction. When using accumulate_outfeed, model callbacks will be called with the same data for the batches which the data was accumulated for. This value is saved/loaded when the model is saved/loaded.

  • gradient_accumulation_reduction_method – (Experimental) Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • pipelining_kwargs – All remaining keyword arguments are forwarded to pipeline(). Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_pipelining_options again.

class keras.ipu.extensions.ModelLayerPipelineStageAssignment(layer, node_index, pipeline_stage=None)

A class to indicate at which pipeline stage a layer in a Model subclass should be executed.

Keras Layers can be called multiple times in order to share weights between layers. Each of these calls produces Tensor output which can be executed in different pipeline stages (as long as these stages are mapped to the same device).

property inbound_layers

Returns the input layers for the layer in this assignment. This can be useful for identifying which specific node_index this is.

property layer

Returns the Keras layer associated with this assignment.

property node_index

Returns the specific call to the layer that produced a Tensor.

property pipeline_stage

Returns the pipeline stage this layer has been assigned to. If None, then this layer has not been assigned to a pipeline stage.

class keras.ipu.extensions.PipelineStage(stage)

A scope within which Keras Layers and/or calls to Keras layers can be assigned to pipeline stages.

Pipeline stages can be assigned to all calls of layer by constructing the layer within a PipelineStage scope as follows:

strategy = ipu.ipu_strategy.IPUStrategy()
input_layer = Input(2)
with strategy.scope():
  with PipelineStage(0):
    x = Dense(4)(input_layer)

  with PipelineStage(1):
    x = Dense(4)(x)

Pipeline stages can also be assigned to individual layer calls, as follows:

strategy = ipu.ipu_strategy.IPUStrategy()
input_layer = Input(2)
l = Dense(4)
with strategy.scope():
  with PipelineStage(0):
    x = l(input_layer)

  with PipelineStage(1):
    x = l(x)

Pipeline stages assigned to layer calls take precedence over those assigned when constructing the layer.

class keras.ipu.extensions.SequentialExtension(*args, **kwargs)
get_pipeline_stage_assignment()

Returns the pipeline stage assignment of the layers in the model.

If set_pipeline_stage_assignment() has been called before, then it returns a copy of the current assignment, otherwise returns a list of SequentialLayerPipelineStageAssignment for each layer in the model in post order (which means that layers are returned in the order they are executed).

print_pipeline_stage_assignment_summary(line_length=None, print_fn=None)

Prints a summary of the pipeline stage assignment of the model.

Parameters

  • line_length – Total length of printed lines (e.g. set this to adapt the display to different terminal window sizes).

  • print_fn – Print function to use. It will be called on each line of the summary. You can set it to a custom function in order to capture the string summary. It defaults to print (prints to stdout).

reset_pipeline_stage_assignment()

Resets the pipeline stage assignment so that the model is no longer pipelined.

set_asynchronous_callbacks(asynchronous=False)

Sets the asynchronous callbacks options when calling fit(), evaluate() and predict().

When running fit(), evaluate() and predict(), the callback functions are called after executing the number of steps specified by steps_per_execution, where each step processes one batch.

Enabling asynchronous callbacks means that the callbacks are invoked after every step, even when steps_per_execution > 1. This can reduce the latency of receiving per-step results and metrics, at the cost of an extra thread running in the background of the application.

Note that this option is ignored for fit() and evaluate() when running a pipelined model and accumulate_outfeed=True (configured via set_pipelining_options).

Parameters

asynchronous – Whether asynchronous callbacks should be enabled.

set_gradient_accumulation_options(gradient_accumulation_steps_per_replica=None, gradient_accumulation_reduction_method='sum', **gradient_accumulation_optimizer_kwargs)

Sets the gradient accumulation options for non-pipelined models which are to be used when training a model.

When set, and gradient_accumulation_steps_per_replica > 1, the optimizer which the current model has been compiled with is wrapped in GradientAccumulationOptimizerV2. This means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when gradient accumulation is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • reduction_method – Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • gradient_accumulation_optimizer_kwargs – All remaining keyword arguments are forwarded to GradientAccumulationOptimizerV2. See the optimizer for all the available arguments. Must not contain opt or num_mini_batches as keys. Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_gradient_accumulation_options again.

set_infeed_queue_options(**kwargs)

Sets the options for all instances of IPUInfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUInfeedQueue is created to efficiently feed data from the dataset to the device. Instances of IPUInfeedQueue can be created with optional arguments, such as prefetch_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUInfeedQueue.

set_outfeed_queue_options(**kwargs)

Sets the options for all instances of IPUOutfeedQueue generated when executing the model.

When using fit(), evalute() and predict(), an instance of IPUOutfeedQueue is created to efficiently feed data from the device to the host. Instances of IPUOutfeedQueue can be created with optional arguments, such as buffer_depth, which can increase the throughput of the model.

Parameters

**kwargs – All keyword arguments are forwarded to IPUOutfeedQueue.

set_pipeline_stage_assignment(pipeline_stage_assignment)

Sets the pipeline stage assignment of all the layers in the model.

Sets the pipeline stage assignment of all the layers in the model which is used to create a model-parallel execution of this Sequential model when calling fit(), evaluate() and predict(). Note that this pipelining stage assignment is ignored when using the call() function on this model.

Parameters

pipeline_stage_assignment – A list of the same length as the number of layers in this model. All elements can be either intergers or instances of SequentialLayerPipelineStageAssignment. If all the elements are integers, then a layer in this model at index i is assigned to a pipeline stage pipeline_stage_assignment[i]. Otherwise, if all the elements are of type SequentialLayerPipelineStageAssignment then a layer in this model at index i is assigned to a pipeline stage indicated by pipeline_stage_assignment[i].pipeline_stage.

Raises

ValueError – pipeline_stage_assignment is not a valid assignment.

set_pipelining_options(gradient_accumulation_steps_per_replica=None, device_mapping=None, accumulate_outfeed=None, gradient_accumulation_reduction_method='sum', **pipelining_kwargs)

Sets the pipelining options, including gradient accumulation options, for pipelined models.

Before training a pipelined model, gradient_accumulation_steps_per_replica argument needs to be set as pipelined models always perform gradient accumulation when training. Setting gradient_accumulation_steps_per_replica > 1 means that each replica will accumulate the gradients for gradient_accumulation_steps_per_replica steps, these accumulated gradients are then all-reduced across the replicas and the weight update is performed.

Gradient Accumulation allows us to simulate bigger batch sizes. For example if we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there is single replica in the system, this simulates an input batch of size 64. If we have a model where each step is of batch size 16 and we set gradient_accumulation_steps_per_replica=4 and there are 4 replicas in the system, this simulates an input batch of size 256.

When training a data-parallel model, enabling gradient accumulation also reduces the communication overhead as the all-reduce of gradients is now performed after each replica has performed gradient_accumulation_steps_per_replica steps instead of after each step.

See the Gradient accumulation section in the documention for more details.

The value of gradient_accumulation_steps_per_replica has no effect when using evaluate() or predict().

Note that the minimize API of the provided optimizer will not be called when pipelining is enabled. As such, overriding minimize in a custom optimizer will cause a ValueError to be raised.

Parameters

  • gradient_accumulation_steps_per_replica – An integer which indicates the number of steps the gradients will be accumulated for in each replica. This value multiplied by the number of replicas needs to divide the steps_per_execution value the model has been compiled with. This value is saved/loaded when the model is saved/loaded.

  • device_mapping – If provided, a list of length equal to the number of pipeline stages assigned in this model. An element at index i in the list represents which IPU the i’th pipeline stage should reside on. This can be used to make sure computational stages which share Keras layers/tf.Variable objects are resident on the same IPU. This value is saved/loaded when the model is saved/loaded.

  • accumulate_outfeed – The metrics from the model are normally enqueued as soon as they’re available. If this option is True, the data will instead be accumulated when they’re available and enqueued at the end of pipeline execution, reducing the amount of host <-> device communication. When used with training, the accumulated metrics are normalised by gradient_accumulation_steps_per_replica. When used with evaluation, the accumulated metrics are normalised by steps_per_epoch. This option is ignored when doing prediction. When using accumulate_outfeed, model callbacks will be called with the same data for the batches which the data was accumulated for. This value is saved/loaded when the model is saved/loaded.

  • gradient_accumulation_reduction_method – (Experimental) Reduction method to use when accumulating gradients. During the iterations in each optimizer step, the computed gradients can either be directly summed up or scaled such that we compute a mean of all gradients for each variable. Computing a mean avoids potential issues with overflow during accumulation especially when using float16, but gives smaller gradients and might require adjusting the learning-rate accordingly. Defaults to GradientAccumulationReductionMethod.SUM (see GradientAccumulationReductionMethod) # pylint: disable=line-too-long

  • pipelining_kwargs – All remaining keyword arguments are forwarded to pipeline(). Note that this dictionary is not serializable, which means that when the model is being saved, these values are not saved. When restoring/loading a model, please call set_pipelining_options again.

class keras.ipu.extensions.SequentialLayerPipelineStageAssignment(layer, pipeline_stage=None)

A class used to indicate which pipeline stage a layer in a Sequential model should be executed in.

property layer

Returns the Keras layer associated with this assignment.

property pipeline_stage

Returns the pipeline stage this layer has been assigned to. If None, then this layer has not been assigned to a pipeline stage.