[Auto Parallel] Improve the APIs

[aoyulong]

PR types

Others

PR changes

APIs

Describe

This pr systematically improve the APIs of auto parallel, and update the codes and unittests as well. These APIs can work in both the dynamic and static graphs, which are transparent for users. The main contributions:

• Engine class (in Paddle/python/paddle/distributed/auto_parallel/engine.py) provides the high-level APIs for distributed training, evaluating and predicting. For example:

  import paddle
  import paddle.vision.transforms as T
  import paddle.distributed.auto_parallel as auto
  from paddle.vision.datasets import MNIST
  
  transform = T.Compose([
      T.Transpose(),
      T.Normalize([127.5], [127.5])
  ])
  train_dataset = MNIST(mode='train', transform=transform)
  valid_dataset = MNIST(mode='test', transform=transform)
  
  model = paddle.vision.models.LeNet()
  loss = paddle.nn.CrossEntropyLoss() 
  optimizer = paddle.optimizer.Adam(
      learning_rate=0.001, parameters=model.parameters())
  metrics = paddle.metric.Accuracy(topk=(1, 2))
  
  engine = auto.Engine(model, loss, optimizer, metrics) 
  # fit 
  engine.fit(train_dataset,
             epochs=2,
             batch_size=64)
  # evaluate 
  engine.evaluate(valid_dataset,
                  batch_size=64)
  # predict
  engine.predict(valid_dataset,
                 batch_size=64)
  # save
  engine.save("./my_model")
  # load 
  engine.load("./my_model")

• shard_tensor and shard_op (in Paddle/python/paddle/distributed/auto_parallel/interface.py) provides the mid-level APIs for users to shard tensors or operators according to their own choices. For example:

  import paddle
  import paddle.distributed.auto_parallel as auto 
  
  mesh = auto.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
  a = paddle.ones([4, 6])
  b = paddle.zeros([4, 6])
  
  # shard_tensor
  auto.shard_tensor(a, mesh, shard_spec=["x", "y"])
  # shard_op, functional style
  auto.shard_op(paddle.add, mesh,
                in_shard_specs=[["x", "y"], ["y", None]],
                out_shard_specs=[[None, "x"]])(a, b)

• Strategy (in Paddle/python/paddle/distributed/auto_parallel/strategy.py) is used to configure the paralleization and optimization behaviors.

• ProcessMesh (in Paddle/python/paddle/distributed/auto_parallel/process_mesh.py) is used to describe the topology of the used processes in the distributed computation.

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[qingqing01]

下面evaluate接口中用的eval_data，这里命名最好保持一致

[aoyulong]

Done.

[JiabinYang]

LGTM

[XieYunshen]

LGTM
单测时间阈值修改
单测迁移