mindspore.parallel.nn.Pipeline

View Source On Gitee
class mindspore.parallel.nn.Pipeline(network, micro_size, stage_config=None, segment_config=None)[source]

Specify the number of micro_batch for pipeline parallelism and the division rules for stage, and specify the segment info.

Note

micro_size must be greater or equal to pipeline stages.

Parameters

  • network (Cell) – The target network to wrap.

  • micro_size (int) – MicroBatch size.

  • stage_config (dict, optional) – Stage configuration for cell's execution in pipeline parallel. Default None.

  • segment_config (dict, optional) – The segment configuration for each cell's execution in pipeline parallel. Default None.

Raises

  • TypeError – The type of net is not cell.

  • TypeError – If the type of micro_size is not int.

  • ValueError – When micro_size is 0 or negative value.

  • KeyError – dict cell name matching exception, there are remaining configuration items after traversing all cell under the current net.

Supported Platforms:

Ascend

Examples

>>> from mindspore.parallel.nn import Pipeline
>>> # Define the network structure of LeNet5. Refer to
>>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
>>> net = LeNet5()
>>> net = Pipeline(net, 4, stage_config={"cell_name_0": 0, "cell_name_1": 1})