Merge pull request #16 from LielinJiang/test-prepare-context

remove init_context

callbacks.py

@@ -16,7 +16,7 @@ import six
 import copy
 
 from progressbar import ProgressBar
-from distributed import get_local_rank
+from paddle.fluid.dygraph.parallel import ParallelEnv
 
 
 def config_callbacks(callbacks=None,
@@ -195,7 +195,7 @@ class ProgBarLogger(Callback):
         self.steps = self.params['steps']
         self.epoch = epoch
         self.train_step = 0
-        if self.verbose and self.epochs and get_local_rank() == 0:
+        if self.verbose and self.epochs and ParallelEnv().local_rank == 0:
             print('Epoch %d/%d' % (epoch + 1, self.epochs))
         self.train_progbar = ProgressBar(num=self.steps, verbose=self.verbose)
 
@@ -213,8 +213,8 @@ class ProgBarLogger(Callback):
         logs = logs or {}
         self.train_step += 1
 
-        if self.train_step % self.log_freq == 0 and self.verbose and get_local_rank(
-        ) == 0:
+        if self.train_step % self.log_freq == 0 and self.verbose and ParallelEnv(
+        ).local_rank == 0:
             # if steps is not None, last step will update in on_epoch_end
             if self.steps and self.train_step < self.steps:
                 self._updates(logs, 'train')
@@ -223,7 +223,7 @@ class ProgBarLogger(Callback):
 
     def on_epoch_end(self, epoch, logs=None):
         logs = logs or {}
-        if self.verbose and get_local_rank() == 0:
+        if self.verbose and ParallelEnv().local_rank == 0:
             self._updates(logs, 'train')
 
     def on_eval_begin(self, logs=None):
@@ -233,7 +233,7 @@ class ProgBarLogger(Callback):
         self.evaled_samples = 0
         self.eval_progbar = ProgressBar(
             num=self.eval_steps, verbose=self.verbose)
-        if get_local_rank() == 0:
+        if ParallelEnv().local_rank == 0:
             print('Eval begin...')
 
     def on_eval_batch_end(self, step, logs=None):
@@ -244,7 +244,7 @@ class ProgBarLogger(Callback):
 
     def on_eval_end(self, logs=None):
         logs = logs or {}
-        if self.verbose and get_local_rank() == 0:
+        if self.verbose and ParallelEnv().local_rank == 0:
             self._updates(logs, 'eval')
             print('Eval samples: %d' % (self.evaled_samples))
 
@@ -258,7 +258,7 @@ class ModelCheckpoint(Callback):
         self.epoch = epoch
 
     def _is_save(self):
-        return self.model and self.save_dir and get_local_rank() == 0
+        return self.model and self.save_dir and ParallelEnv().local_rank == 0
 
     def on_epoch_end(self, epoch, logs=None):
         if self._is_save() and self.epoch % self.save_freq == 0:

distributed.py

@@ -13,30 +13,21 @@
 # limitations under the License.
 import os
 import sys
+import six
 import time
 import math
 import socket
 import contextlib
-from contextlib import closing
-from six import string_types
 import numpy as np
-from collections import OrderedDict
-from paddle import fluid
-import paddle.fluid.unique_name as nameGen
-from paddle.fluid import core
 
-from paddle.fluid import framework
+from paddle import fluid
 from paddle.fluid.layers import collective
-from paddle.fluid.dygraph import to_variable, no_grad, layers
-from paddle.fluid.framework import Variable
-from paddle.fluid.executor import global_scope
+from paddle.fluid.dygraph.parallel import ParallelEnv, ParallelStrategy
+from paddle.fluid.io import BatchSampler
+
 
-from paddle.fluid.dygraph.parallel import Env, DataParallel, ParallelStrategy
-from paddle.fluid.layers.collective import _c_allreduce, _c_allgather, _c_broadcast, \
-                                            _c_sync_comm_stream, _c_sync_calc_stream
-from paddle.fluid.io import BatchSampler, DataLoader
+_parallel_context_initialized = False
 
-__parallel_context_init = False
 
 class DistributedBatchSampler(BatchSampler):
     """Sampler that restricts data loading to a subset of the dataset.
@@ -71,9 +62,10 @@ class DistributedBatchSampler(BatchSampler):
         self.shuffle = shuffle
         assert isinstance(drop_last, bool), \
                 "drop_last should be a boolean number"
+
         self.drop_last = drop_last
-        self.nranks = get_nranks()
-        self.local_rank = get_local_rank()
+        self.nranks = ParallelEnv().nranks
+        self.local_rank = ParallelEnv().local_rank
         self.epoch = 0
         self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.nranks))
         self.total_size = self.num_samples * self.nranks
@@ -106,27 +98,22 @@ class DistributedBatchSampler(BatchSampler):
         num_samples += int(not self.drop_last) * (self.batch_size - 1)
         return num_samples // self.batch_size
 
-
-def _all_gather(x, nranks, ring_id=0, use_calc_stream=True):
-    return _c_allgather(x, nranks, ring_id=ring_id, use_calc_stream=use_calc_stream)
+    def set_epoch(self, epoch):
+        self.epoch = epoch
 
 
-def get_local_rank():
-    return Env().local_rank
-
-
-def get_nranks():
-    return Env().nranks
+def _all_gather(x, nranks, ring_id=0, use_calc_stream=True):
+    return collective._c_allgather(x, nranks, ring_id=ring_id, use_calc_stream=use_calc_stream)
 
 
 def wait_server_ready(endpoints):
-    assert not isinstance(endpoints, string_types)
+    assert not isinstance(endpoints, six.string_types)
     while True:
         all_ok = True
         not_ready_endpoints = []
         for ep in endpoints:
             ip_port = ep.split(":")
-            with closing(
+            with contextlib.closing(
                     socket.socket(socket.AF_INET,
                                   socket.SOCK_STREAM)) as sock:
                 sock.settimeout(2)
@@ -135,10 +122,6 @@ def wait_server_ready(endpoints):
                     all_ok = False
                     not_ready_endpoints.append(ep)
         if not all_ok:
-            sys.stderr.write("server not ready, wait 3 sec to retry...\n")
-            sys.stderr.write("not ready endpoints:" + str(
-                not_ready_endpoints) + "\n")
-            sys.stderr.flush()
             time.sleep(3)
         else:
             break
@@ -154,9 +137,9 @@ def init_communicator(program, rank, nranks, wait_port,
         wait_server_ready(other_endpoints)
     block = program.global_block()
     nccl_id_var = block.create_var(
-        name=nameGen.generate('nccl_id'),
+        name=fluid.unique_name.generate('nccl_id'),
         persistable=True,
-        type=core.VarDesc.VarType.RAW)
+        type=fluid.core.VarDesc.VarType.RAW)
 
     block.append_op(
         type='c_gen_nccl_id',
@@ -181,23 +164,23 @@ def init_communicator(program, rank, nranks, wait_port,
 
 def prepare_distributed_context(place=None):
     if place is None:
-        place = fluid.CUDAPlace(Env().dev_id) if Env().nranks > 1 \
+        place = fluid.CUDAPlace(ParallelEnv().dev_id) if ParallelEnv().nranks > 1 \
             else fluid.CUDAPlace(0)
     
     strategy = ParallelStrategy()
-    strategy.nranks = Env().nranks
-    strategy.local_rank = Env().local_rank
-    strategy.trainer_endpoints = Env().trainer_endpoints
-    strategy.current_endpoint = Env().current_endpoint
+    strategy.nranks = ParallelEnv().nranks
+    strategy.local_rank = ParallelEnv().local_rank
+    strategy.trainer_endpoints = ParallelEnv().trainer_endpoints
+    strategy.current_endpoint = ParallelEnv().current_endpoint
 
     if strategy.nranks < 2:
         return
 
-    global __parallel_context_init
+    global _parallel_context_initialized
 
-    if not __parallel_context_init and isinstance(place, core.CUDAPlace):
+    if not _parallel_context_initialized and isinstance(place, fluid.CUDAPlace):
         def _init_context():
-            communicator_prog = framework.Program()
+            communicator_prog = fluid.Program()
             init_communicator(communicator_prog, strategy.local_rank, strategy.nranks, 
                             True, strategy.current_endpoint, strategy.trainer_endpoints)
             exe = fluid.Executor(place)
@@ -213,57 +196,5 @@ def prepare_distributed_context(place=None):
     else:
         assert ("Only support CUDAPlace for now.")
 
-    __parallel_context_init = True
+    _parallel_context_initialized = True
     return strategy
\ No newline at end of file
-
-
-class DistributedDataParallel(DataParallel):
-    def __init__(self, layers, strategy=None):
-        if strategy is None:
-            strategy = ParallelStrategy()
-            strategy.nranks = Env().nranks
-            strategy.local_rank = Env().local_rank
-            strategy.trainer_endpoints = Env().trainer_endpoints
-            strategy.current_endpoint = Env().current_endpoint
-
-        super(DistributedDataParallel, self).__init__(layers, strategy)
-
-    @no_grad
-    def apply_collective_grads(self):
-        """
-        AllReduce the Parameters' gradient.
-        """
-        if not self._is_data_parallel_mode():
-            return
-
-        grad_var_set = set()
-        grad_vars = []
-        for param in self._layers.parameters():
-            # NOTE(zcd): The grad_ivar maybe no generated.
-            if param.trainable and param._grad_ivar():
-                g_var = param._grad_ivar()
-                grad_vars.append(g_var)
-                assert g_var not in grad_var_set
-                grad_var_set.add(g_var)
-
-        mega_bytes = 128 * 1024 * 1024
-        group_idx = 0
-        memory_counter = 0
-        grad_var_groups = OrderedDict()
-        dtype = grad_vars[0].dtype
-        for g_var in grad_vars:
-            # Note: the dtype of the same group should be the same.
-            bytes = np.prod(g_var.shape) * core.size_of_dtype(g_var.dtype)
-            if memory_counter < mega_bytes and dtype == g_var.dtype:
-                memory_counter += bytes
-            else:
-                memory_counter = bytes
-                group_idx += 1
-            grad_var_groups.setdefault(group_idx, []).append(g_var)
-
-        coalesced_grads_and_vars = self._coalesce_tensors(grad_var_groups)
-
-        for coalesced_grad, _, _ in coalesced_grads_and_vars:
-            collective._c_allreduce(coalesced_grad, coalesced_grad, use_calc_stream=True)
-
-        self._split_tensors(coalesced_grads_and_vars)

mnist.py

@@ -26,7 +26,7 @@ from paddle.fluid.optimizer import Momentum
 from paddle.fluid.dygraph.nn import Conv2D, Pool2D, Linear
 from paddle.fluid.io import MNIST as MnistDataset
 
-from model import Model, CrossEntropy, Input, init_context
+from model import Model, CrossEntropy, Input, set_device
 from metrics import Accuracy
 
 
@@ -106,7 +106,8 @@ class MNIST(Model):
 
 
 def main():
-    init_context('dynamic' if FLAGS.dynamic else 'static')
+    device = set_device(FLAGS.device)
+    fluid.enable_dygraph(device) if FLAGS.dynamic else None
 
     train_dataset = MnistDataset(mode='train')
     val_dataset = MnistDataset(mode='test')
@@ -118,7 +119,13 @@ def main():
     optim = Momentum(
         learning_rate=FLAGS.lr, momentum=.9, parameter_list=model.parameters())
 
-    model.prepare(optim, CrossEntropy(), Accuracy(topk=(1, 2)), inputs, labels)
+    model.prepare(
+        optim,
+        CrossEntropy(),
+        Accuracy(topk=(1, 2)),
+        inputs,
+        labels,
+        device=FLAGS.device)
     if FLAGS.resume is not None:
         model.load(FLAGS.resume)
 
@@ -131,6 +138,8 @@ def main():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser("CNN training on MNIST")
+    parser.add_argument(
+        "--device", type=str, default='gpu', help="device to use, gpu or cpu")
     parser.add_argument(
         "-d", "--dynamic", action='store_true', help="enable dygraph mode")
     parser.add_argument(

model.py

@@ -20,25 +20,34 @@ import pickle
 import numpy as np
 import six
 import warnings
-from collections import Iterable
-from collections import OrderedDict
 
-from collections import OrderedDict
+from collections import Iterable
 from paddle import fluid
 from paddle.fluid.framework import in_dygraph_mode, Variable
 from paddle.fluid.executor import global_scope
 from paddle.fluid.io import is_belong_to_optimizer
 from paddle.fluid.dygraph.base import to_variable
-
+from paddle.fluid.dygraph.parallel import ParallelEnv
 from paddle.fluid.incubate.fleet.collective import fleet, DistributedStrategy
-import paddle.fluid.incubate.fleet.base.role_maker as role_maker
-import distributed
-from distributed import DistributedBatchSampler
-from paddle.fluid.io import DataLoader
+from paddle.fluid.incubate.fleet.base import role_maker
+from paddle.fluid.io import DataLoader, Dataset
+
+from distributed import DistributedBatchSampler, _all_gather, prepare_distributed_context, _parallel_context_initialized
 from metrics import Metric
 from callbacks import config_callbacks
 
-__all__ = ['Model', 'Loss', 'CrossEntropy', 'Input']
+__all__ = ['Model', 'Loss', 'CrossEntropy', 'Input', 'set_device']
+
+
+def set_device(device):
+    assert isinstance(device, six.string_types) and device.lower() in ['cpu', 'gpu'], \
+    "Expected device in ['cpu', 'gpu'], but got {}".format(device)
+
+    place = fluid.CUDAPlace(ParallelEnv().dev_id) \
+            if device.lower() == 'gpu' and fluid.is_compiled_with_cuda() \
+                else fluid.CPUPlace()
+
+    return place
 
 
 def to_list(value):
@@ -84,18 +93,6 @@ def extract_args(func):
         return inspect.getargspec(func)[0]
 
 
-def init_context(backend):
-    assert isinstance(backend, str) and backend.lower() in ['dynamic', 'static'], \
-        "Expected backend in ['dynamic', 'static'], but got {}".format(backend)
-
-    place = fluid.CUDAPlace(distributed.Env().dev_id) if \
-                distributed.Env().nranks > 1 else fluid.CUDAPlace(0)
-    distributed.prepare_distributed_context(place)
-    backend = backend.lower()
-    if backend == 'dynamic':
-        fluid.enable_dygraph(place)
-
-
 class Input(fluid.dygraph.Layer):
     def __init__(self, shape=None, dtype=None, name=None):
         super(Input, self).__init__()
@@ -161,8 +158,8 @@ class StaticGraphAdapter(object):
             'test_batch': 0
         }
 
-        self._nranks = distributed.Env().nranks
-        self._local_rank = distributed.Env().local_rank
+        self._nranks = ParallelEnv().nranks
+        self._local_rank = ParallelEnv().local_rank
 
     @property
     def mode(self):
@@ -267,7 +264,8 @@ class StaticGraphAdapter(object):
                 # When using static learning rate, static-graph would make it
                 # a persistable var named 'unique_name.generate("learning_rate")',
                 # However, dygraph wouldn't save it.
-                if var.name not in state: continue
+                if var.name not in state:
+                    continue
             else:
                 # moment and other accumulators
                 if var.name not in converted_state:
@@ -416,8 +414,10 @@ class StaticGraphAdapter(object):
         metrics = []
         with fluid.program_guard(prog, self._startup_prog):
             if isinstance(self.model._inputs, dict):
-                ins = [self.model._inputs[n] \
-                    for n in extract_args(self.model.forward) if n != 'self']
+                ins = [
+                    self.model._inputs[n]
+                    for n in extract_args(self.model.forward) if n != 'self'
+                ]
             else:
                 ins = self.model._inputs
             lbls = self.model._labels if self.model._labels else []
@@ -430,14 +430,9 @@ class StaticGraphAdapter(object):
                 losses = self.model._loss_function(outputs, labels)
 
             if self._nranks > 1 and mode != 'train':
-                outputs = [
-                    distributed._all_gather(o, self._nranks) for o in outputs
-                ]
+                outputs = [_all_gather(o, self._nranks) for o in outputs]
                 if mode != 'test':
-                    labels = [
-                        distributed._all_gather(l, self._nranks)
-                        for l in labels
-                    ]
+                    labels = [_all_gather(l, self._nranks) for l in labels]
 
             if mode != 'test':
                 for metric in self.model._metrics:
@@ -474,24 +469,15 @@ class StaticGraphAdapter(object):
         if compiled_prog is not None:
             return compiled_prog
 
-        device = self.model._device
-        device_ids = self.model._device_ids
+        assert self.model._place is not None, \
+            "device is not set, please call `model.prepare()` first"
 
-        if device.lower() == 'gpu':
-            places = fluid.cuda_places(device_ids)
-        else:
-            places = fluid.cpu_places(len(device_ids) if device_ids else None)
+        place = self.model._place
 
         # XXX *ALL WEIGHTS* should be initialized upon model construction
         # even if `forward()` may run different code path for different mode
         # therefore startup program only needs to run once
         if self._executor is None:
-            if self._nranks > 1 and device.lower() == 'gpu':
-                gpu_id = int(distributed.Env().dev_id)
-                place = fluid.CUDAPlace(gpu_id) if device.lower(
-                ) == 'gpu' else fluid.CPUPlace()
-            else:
-                place = places[0]
             self._executor = fluid.Executor(place)
             # XXX incremental initialization
             uninitialized = []
@@ -509,14 +495,8 @@ class StaticGraphAdapter(object):
         if self._nranks < 2:
             compiled_prog = fluid.CompiledProgram(prog)
         else:
-            compiled_prog = prog  #fleet.main_program
-
-        if len(places) > 1:
-            loss_name = None
-            if mode == 'train' and self._loss_endpoint is not None:
-                loss_name = self._loss_endpoint.name
-            compiled_prog = compiled_prog.with_data_parallel(
-                loss_name=loss_name, places=places)
+            compiled_prog = prog
+
         self._compiled_progs[mode] = compiled_prog
 
 
@@ -524,8 +504,8 @@ class DynamicGraphAdapter(object):
     def __init__(self, model):
         super(DynamicGraphAdapter, self).__init__()
         self.model = model
-        self._nranks = distributed.Env().nranks
-        self._local_rank = distributed.Env().local_rank
+        self._nranks = ParallelEnv().nranks
+        self._local_rank = ParallelEnv().local_rank
         self._merge_count = {
             'eval_total': 0,
             'test_total': 0,
@@ -534,7 +514,13 @@ class DynamicGraphAdapter(object):
         }
 
         if self._nranks > 1:
-            self.ddp_model = distributed.DistributedDataParallel(self.model)
+            stradegy = fluid.dygraph.parallel.ParallelStrategy()
+            stradegy.nranks = ParallelEnv().nranks
+            stradegy.local_rank = ParallelEnv().local_rank
+            stradegy.trainer_endpoints = ParallelEnv().trainer_endpoints
+            stradegy.current_endpoint = ParallelEnv().current_endpoint
+            self.ddp_model = fluid.dygraph.parallel.DataParallel(self.model,
+                                                                 stradegy)
 
     @property
     def mode(self):
@@ -590,11 +576,8 @@ class DynamicGraphAdapter(object):
         else:
             losses = []
         if self._nranks > 1:
-            outputs = [
-                distributed._all_gather(o, self._nranks)
-                for o in to_list(outputs)
-            ]
-            labels = [distributed._all_gather(l, self._nranks) for l in labels]
+            outputs = [_all_gather(o, self._nranks) for o in to_list(outputs)]
+            labels = [_all_gather(l, self._nranks) for l in labels]
         metrics = []
         for metric in self.model._metrics:
             # cut off padding value.
@@ -630,10 +613,7 @@ class DynamicGraphAdapter(object):
         inputs = [to_variable(x) for x in to_list(inputs)]
         outputs = self.model.forward(*inputs)
         if self._nranks > 2:
-            outputs = [
-                distributed._all_gather(o, self._nranks)
-                for o in to_list(outputs)
-            ]
+            outputs = [_all_gather(o, self._nranks) for o in to_list(outputs)]
         return [to_numpy(o) for o in to_list(outputs)]
 
     def parameters(self, *args, **kwargs):
@@ -720,10 +700,6 @@ class Model(fluid.dygraph.Layer):
         self._optimizer = None
         self._test_dataloader = None
 
-        # init multiple gpus context
-        self._place = fluid.CUDAPlace(distributed.Env().dev_id) \
-                    if distributed.Env().nranks > 1 else fluid.CUDAPlace(0)
-
         # init backend
         if fluid.in_dygraph_mode():
             self._adapter = DynamicGraphAdapter(self)
@@ -740,7 +716,7 @@ class Model(fluid.dygraph.Layer):
         return self._adapter.test(*args, **kwargs)
 
     def save(self, *args, **kwargs):
-        if distributed.get_local_rank() == 0:
+        if ParallelEnv().local_rank == 0:
             return self._adapter.save(*args, **kwargs)
 
     def load(self, path, skip_mismatch=False, reset_optimizer=False):
@@ -855,6 +831,35 @@ class Model(fluid.dygraph.Layer):
                 The default is None.
         """
 
+        if isinstance(device, fluid.CUDAPlace) or \
+            (isinstance(device, six.string_types) and device.lower() == 'gpu') \
+            or (device is None and fluid.is_compiled_with_cuda()):
+            if isinstance(device, fluid.CUDAPlace):
+                self._place = device
+            else:
+                self._place = fluid.CUDAPlace(ParallelEnv().dev_id) \
+                    if ParallelEnv().nranks > 1 else fluid.CUDAPlace(0)
+
+            global _parallel_context_initialized
+            if ParallelEnv().nranks > 1 and not _parallel_context_initialized:
+                if fluid.in_dygraph_mode():
+                    fluid.disable_dygraph()
+                    fluid.enable_dygraph(self._place)
+                    fluid.dygraph.parallel.prepare_context()
+                else:
+                    prepare_distributed_context(self._place)
+
+                _parallel_context_initialized = True
+        elif isinstance(device, fluid.CPUPlace):
+            self._place = device
+        elif (isinstance(device, six.string_types) and device.lower() == 'cpu') \
+            or (device is None):
+            self._place = fluid.CPUPlace()
+        else:
+            raise ValueError(
+                "Expected device in ('gpu', 'cpu', fluid.CUDAPlace, fluid.CPUPlace, None), \
+                but got {}".format(device))
+
         self._optimizer = optimizer
         if loss_function:
             if not isinstance(loss_function, Loss):
@@ -871,25 +876,20 @@ class Model(fluid.dygraph.Layer):
         metrics = metrics or []
         for metric in to_list(metrics):
             assert isinstance(metric, Metric), \
-                "{} is not sub class of Metric".format(metric.__class__.__name__)
+                "{} is not sub class of Metric".format(
+                    metric.__class__.__name__)
         self._metrics = to_list(metrics)
 
         self._inputs = inputs
         self._labels = labels
-        self._device = device
 
-        if device is None:
-            self._device = 'GPU' if fluid.is_compiled_with_cuda() else 'CPU'
-        self._device_ids = device_ids
         if not in_dygraph_mode():
             self._adapter.prepare()
 
     def fit(
             self,
-            train_dataset=None,
-            eval_dataset=None,
-            train_loader=None,
-            eval_loader=None,
+            train_data=None,
+            eval_data=None,
             batch_size=1,
             epochs=1,
             eval_freq=1,
@@ -904,9 +904,16 @@ class Model(fluid.dygraph.Layer):
         """
         FIXME: add more comments and usage
         Args:
-            train_loader (DataLoader): An iterable data loader is used for train.
-            eval_loader (DataLoader): An iterable data loader is used for
+            train_data (Dataset|DataLoader): An iterable data loader is used for 
+                train. An instance of paddle.fluid.io.Dataset or 
+                paddle.fluid.io.Dataloader is recomended.
+            eval_data (Dataset|DataLoader): An iterable data loader is used for
                 evaluation at the end of epoch. If None, will not do evaluation. 
+                An instance of paddle.fluid.io.Dataset or paddle.fluid.io.Dataloader 
+                is recomended.
+            batch_size (int): Integer number. The batch size of train_data and eval_data. 
+                When train_data and eval_data are both the instance of Dataloader, this 
+                parameter will be ignored.
             epochs (int): Integer number. The number of epochs to train the model.
             eval_freq (int): The frequency, in number of epochs, an evalutation
                 is performed.
@@ -917,47 +924,57 @@ class Model(fluid.dygraph.Layer):
             save_freq (int): The frequency, in number of epochs, to save checkpoint.
             verbose (int): The verbosity mode, should be 0, 1, or 2.
                 0 = silent, 1 = progress bar, 2 = one line per epoch.
+            drop_last (bool): whether drop the last incomplete batch of train_data 
+                when dataset size is not divisible by the batch size. When train_data 
+                is an instance of Dataloader, this parameter will be ignored.
+            shuffle (bool): whther to shuffle train_data. When train_data is an instance 
+                of Dataloader, this parameter will be ignored.
+            num_workers (int): the number of subprocess to load data, 0 for no subprocess 
+                used and loading data in main process. When train_data and eval_data are
+                both the instance of Dataloader, this parameter will be ignored.
             callbacks (Callback|None): A list of `Callback` instances to apply
                 during training. If None, `ProgBarLogger` and `ModelCheckpoint`
                 are automatically inserted.
         """
 
-        assert train_dataset is not None or train_loader is not None, \
-                "train_dataset or train_loader must be given"
-
-        assert (train_loader is not None and train_dataset is None) or \
-               (train_loader is None and train_dataset is not None), \
-                "train_dataset should not be set when train_loader is given"
+        assert train_data is not None, \
+                "train_data must be given!"
 
         if fluid.in_dygraph_mode():
             feed_list = None
         else:
             feed_list = [x.forward() for x in self._inputs + self._labels]
 
-        if train_loader is None:
+        if isinstance(train_data, Dataset):
             train_sampler = DistributedBatchSampler(
-                train_dataset,
+                train_data,
                 batch_size=batch_size,
                 shuffle=shuffle,
                 drop_last=drop_last)
             train_loader = DataLoader(
-                train_dataset,
+                train_data,
                 batch_sampler=train_sampler,
                 places=self._place,
                 feed_list=feed_list,
                 num_workers=num_workers,
                 return_list=True)
+        else:
+            train_loader = train_data
 
-        if eval_loader is None and eval_dataset is not None:
+        if eval_data is not None and isinstance(eval_data, Dataset):
             eval_sampler = DistributedBatchSampler(
-                eval_dataset, batch_size=batch_size)
+                eval_data, batch_size=batch_size)
             eval_loader = DataLoader(
-                eval_dataset,
+                eval_data,
                 batch_sampler=eval_sampler,
                 places=self._place,
                 feed_list=feed_list,
                 num_workers=num_workers,
                 return_list=True)
+        elif eval_data is not None:
+            eval_loader = eval_data
+        else:
+            eval_loader = None
 
         do_eval = eval_loader is not None
         self._test_dataloader = eval_loader
@@ -1010,7 +1027,7 @@ class Model(fluid.dygraph.Layer):
                 logs['step'] = step
                 if mode == 'train' or self._adapter._merge_count.get(
                         mode + '_batch', 0) <= 0:
-                    logs['batch_size'] = batch_size * distributed.Env().nranks
+                    logs['batch_size'] = batch_size * ParallelEnv().nranks
                 else:
                     logs['batch_size'] = self._adapter._merge_count[mode +
                                                                     '_batch']
@@ -1035,7 +1052,7 @@ class Model(fluid.dygraph.Layer):
                 loader = eval_loader
                 if not isinstance(eval_loader, Iterable):
                     loader = eval_loader()
-                logs = _run_one_epoch(eval_loader(), cbks, 'eval')
+                logs = _run_one_epoch(loader, cbks, 'eval')
                 cbks.on_end('eval', logs)
 
         cbks.on_end('train', logs)

tests/test_model.py

@@ -28,7 +28,7 @@ import contextlib
 import paddle
 from paddle import fluid
 from paddle.fluid.dygraph.nn import Conv2D, Pool2D, Linear
-from model import Model, CrossEntropy, Input, Loss, init_context
+from model import Model, CrossEntropy, Input, Loss, set_device
 from metrics import Accuracy
 from callbacks import ProgBarLogger
 from paddle.fluid.io import BatchSampler, DataLoader
@@ -141,7 +141,8 @@ class MyCrossEntropy(Loss):
 
 class TestModel(unittest.TestCase):
     def fit(self, dynamic, is_mlp=False):
-        init_context('dynamic' if dynamic else 'static')
+        device = set_device('gpu')
+        fluid.enable_dygraph(device) if dynamic else None
 
         im_shape = (-1, 784)
         batch_size = 128
@@ -156,7 +157,7 @@ class TestModel(unittest.TestCase):
         optim = fluid.optimizer.Momentum(
             learning_rate=0.01, momentum=.9, parameter_list=model.parameters())
         loss = CrossEntropy() if not is_mlp else MyCrossEntropy()
-        model.prepare(optim, loss, Accuracy(), inputs, labels)
+        model.prepare(optim, loss, Accuracy(), inputs, labels, device=device)
         cbk = ProgBarLogger(50)
         model.fit(train_dataset,
                   val_dataset,