[Unstructured Pruner]add support for cpu training (#898)

demo/dygraph/unstructured_pruning/train.py

@@ -22,6 +22,7 @@ _logger = get_logger(__name__, level=logging.INFO)
 parser = argparse.ArgumentParser(description=__doc__)
 add_arg = functools.partial(add_arguments, argparser=parser)
 # yapf: disable
+add_arg('use_gpu',          bool, True,               "Whether to use GPU for training or not. Default: True")
 add_arg('batch_size',       int,  64,                 "Minibatch size. Default: 64")
 add_arg('batch_size_for_validation',       int,  64,                 "Minibatch size for validation. Default: 64")
 add_arg('lr',               float,  0.05,               "The learning rate used to fine-tune pruned model. Default: 0.05")
@@ -105,7 +106,10 @@ def create_unstructured_pruner(model, args, configs=None):
 
 
 def compress(args):
+    if args.use_gpu:
         place = paddle.set_device('gpu')
+    else:
+        place = paddle.set_device('cpu')
 
     trainer_num = paddle.distributed.get_world_size()
     use_data_parallel = trainer_num != 1

demo/dygraph/unstructured_pruning/train_imagenet_cpu_base.sh

+#!/bin/bash  
+CUDA_VISIBLE_DEVICES='' python \
+          train.py \
+          --batch_size 64 \
+          --data imagenet \
+          --pruning_mode ratio \
+          --ratio 0.55 \
+          --lr 0.05 \
+          --use_gpu False

demo/unstructured_prune/train.py

@@ -21,6 +21,7 @@ _logger = get_logger(__name__, level=logging.INFO)
 parser = argparse.ArgumentParser(description=__doc__)
 add_arg = functools.partial(add_arguments, argparser=parser)
 # yapf: disable
+add_arg('use_gpu',          bool, True,               "Whether to use gpu for traning or not. Defauly: True")
 add_arg('batch_size',       int,  64,                 "Minibatch size. Default: 64")
 add_arg('batch_size_for_validation',       int,  64,                 "Minibatch size for validation. Default: 64")
 add_arg('model',            str,  "MobileNet",                "The target model.")
@@ -137,7 +138,10 @@ def compress(args):
     image_shape = [int(m) for m in image_shape.split(",")]
     assert args.model in model_list, "{} is not in lists: {}".format(args.model,
                                                                      model_list)
+    if args.use_gpu:
         places = paddle.static.cuda_places()
+    else:
+        places = paddle.static.cpu_places()
     place = places[0]
     exe = paddle.static.Executor(place)
 

demo/unstructured_prune/train_imagenet_cpu_base.sh

+CUDA_VISIBLE_DEVICES='' python train.py \
+          --batch_size 64 \
+          --data imagenet \
+          --pruning_mode ratio \
+          --ratio 0.55 \
+          --lr 0.05 \
+          --model MobileNet \
+          --pretrained_model "MobileNetV1_pretrained" \
+          --use_gpu False \