[HybridParallel] Add pipeline layer in dygraph

[ForFishes]

PR types

New features

PR changes

APIs

Describe

[HybridParallel] Add pipeline layer in dygraph
支持使用LayerDesc描述动态图组网，LayerDesc描述的layer不会被初始化，从而节省内存开销。支持uniform切图。使用示例如下：

import numpy as np
import os
import paddle
from paddle.distributed import fleet
import copy
from paddle.fluid.dygraph.container import Sequential
import paddle.nn as nn
from paddle.fluid.dygraph.layers import Layer
from paddle.distributed.fleet.meta_parallel import LayerDesc, PipelineLayer
import paddle.nn.functional as F

class AlexNetPipeDesc(PipelineLayer):
    def __init__(self, num_classes=10, **kwargs):
        self.num_classes = num_classes
        decs = [
            LayerDesc(
                nn.Conv2D, 3, 64, kernel_size=11, stride=4, padding=5),
            LayerDesc(nn.ReLU),
            LayerDesc(
                nn.MaxPool2D, kernel_size=2, stride=2),
            LayerDesc(
                nn.Conv2D, 64, 192, kernel_size=5, padding=2),
            F.relu,
            LayerDesc(
                nn.MaxPool2D, kernel_size=2, stride=2),
            LayerDesc(
                nn.Conv2D, 192, 384, kernel_size=3, padding=1),
            F.relu,
            LayerDesc(
                nn.Conv2D, 384, 256, kernel_size=3, padding=1),
            F.relu,
            LayerDesc(
                nn.Conv2D, 256, 256, kernel_size=3, padding=1),
            F.relu,
            LayerDesc(
                nn.MaxPool2D, kernel_size=2, stride=2),
            lambda x: x.flatten(),
            LayerDesc(nn.Linear, 256, self.num_classes),  # classifier
        ]
        super(AlexNetPipeDesc, self).__init__(
            layers=decs, loss_fn=nn.CrossEntropyLoss(), **kwargs)

通过PipelineLyaer，会打印相应的切图信息。如果切分2层，

start segment network..
stage=0, global_rank=1 ,layer_number=7
0: Conv2D(3, 64, kernel_size=[11, 11], stride=[4, 4], padding=5, data_format=NCHW)
1: ReLU()
2: MaxPool2D(kernel_size=2, stride=2, padding=0)
3: Conv2D(64, 192, kernel_size=[5, 5], padding=2, data_format=NCHW)
4: ReLU()
5: MaxPool2D(kernel_size=2, stride=2, padding=0)
6: Conv2D(192, 384, kernel_size=[3, 3], padding=1, data_format=NCHW)

stage=1, global_rank=1 ,layer_number=8
7: ReLU()
8: Conv2D(384, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
9: ReLU()
10: Conv2D(256, 256, kernel_size=[3, 3], padding=1, data_format=NCHW)
11: ReLU()
12: MaxPool2D(kernel_size=2, stride=2, padding=0)
13: <function AlexNetPipe.to_layers.<locals>.<lambda> at 0x7f39200daae8>
14: Linear(in_features=256, out_features=10, dtype=float32)
loss: CrossEntropyLoss

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.