diff --git a/mnist.py b/mnist.py index 26d78246ad5b431a7b7458c540f92641f801155c..745dc2f06e54136756ce5ae4f3b077c24468dd1d 100644 --- a/mnist.py +++ b/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 +from model import Model, CrossEntropy, Input, set_device from metrics import Accuracy @@ -106,7 +106,8 @@ class MNIST(Model): def main(): - fluid.enable_dygraph() if FLAGS.dynamic else None + 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( diff --git a/model.py b/model.py index 6f4f3b44763f166a4b0afe48a86219e6e5fce788..13ccd57ef9c8884342a519a2e1827ece0c5e3be0 100644 --- a/model.py +++ b/model.py @@ -36,7 +36,18 @@ from distributed import DistributedBatchSampler, _all_gather, prepare_distribute 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.CUDAPlace(0) + + return place def to_list(value): @@ -458,24 +469,21 @@ class StaticGraphAdapter(object): if compiled_prog is not None: return compiled_prog - device = self.model._device - device_ids = self.model._device_ids - - if device.lower() == 'gpu': - places = fluid.cuda_places(device_ids) - else: - places = fluid.cpu_places(len(device_ids) if device_ids else None) + assert self.model._place is not None, \ + "model not ready, please call `model.prepare()` first" + place = self.model._place + fluid.cpu_places # 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(ParallelEnv().dev_id) - place = fluid.CUDAPlace(gpu_id) if device.lower( - ) == 'gpu' else fluid.CPUPlace() - else: - place = places[0] + # if self._nranks > 1 and device.lower() == 'gpu': + # gpu_id = int(ParallelEnv().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 = [] @@ -495,12 +503,6 @@ class StaticGraphAdapter(object): else: compiled_prog = prog - 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) self._compiled_progs[mode] = compiled_prog @@ -704,21 +706,6 @@ class Model(fluid.dygraph.Layer): self._optimizer = None self._test_dataloader = None - # init multiple gpus context - 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 - # init backend if fluid.in_dygraph_mode(): self._adapter = DynamicGraphAdapter(self) @@ -850,6 +837,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): @@ -872,11 +888,7 @@ class Model(fluid.dygraph.Layer): 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() @@ -1046,7 +1058,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) diff --git a/tests/test_model.py b/tests/test_model.py index 8ca19a884494264b81977b78900093fad8f9b76f..87c9e5731f4f73d93b0a7dda70b771f76fdabae3 100644 --- a/tests/test_model.py +++ b/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 +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): - fluid.enable_dygraph() if dynamic else None + 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,