Merge branch 'develop' into 'develop'

Add distillation demo_guide link

See merge request !69

demo/distillation/train.py → demo/distillation/distillation_demo.py

@@ -13,8 +13,7 @@ import numpy as np
 import paddle.fluid as fluid
 sys.path.append(sys.path[0] + "/../")
 import models
-import imagenet_reader as reader
-from utility import add_arguments, print_arguments
+from utility import add_arguments, print_arguments, _download, _decompress
 from paddleslim.dist import merge, l2_loss, soft_label_loss, fsp_loss
 
 logging.basicConfig(format='%(asctime)s-%(levelname)s: %(message)s')
@@ -33,12 +32,12 @@ add_arg('lr_strategy',      str,  "piecewise_decay",   "The learning rate decay
 add_arg('l2_decay',         float,  3e-5,               "The l2_decay parameter.")
 add_arg('momentum_rate',    float,  0.9,               "The value of momentum_rate.")
 add_arg('num_epochs',       int,  120,               "The number of total epochs.")
-add_arg('data',             str, "mnist",                 "Which data to use. 'mnist' or 'imagenet'")
+add_arg('data',             str, "cifar10",                 "Which data to use. 'cifar10' or 'imagenet'")
 add_arg('log_period',       int, 20,                 "Log period in batches.")
 add_arg('model',            str,  "MobileNet",          "Set the network to use.")
 add_arg('pretrained_model', str,  None,                "Whether to use pretrained model.")
 add_arg('teacher_model',    str,  "ResNet50",          "Set the teacher network to use.")
-add_arg('teacher_pretrained_model', str,  "../pretrain/ResNet50_pretrained",                "Whether to use pretrained model.")
+add_arg('teacher_pretrained_model', str,  "./ResNet50_pretrained",                "Whether to use pretrained model.")
 parser.add_argument('--step_epochs', nargs='+', type=int, default=[30, 60, 90], help="piecewise decay step")
 # yapf: enable
 
@@ -76,12 +75,12 @@ def create_optimizer(args):
 
 
 def compress(args):
-    if args.data == "mnist":
-        import paddle.dataset.mnist as reader
-        train_reader = reader.train()
-        val_reader = reader.test()
+    if args.data == "cifar10":
+        import paddle.dataset.cifar as reader
+        train_reader = reader.train10()
+        val_reader = reader.test10()
         class_dim = 10
-        image_shape = "1,28,28"
+        image_shape = "3,32,32"
     elif args.data == "imagenet":
         import imagenet_reader as reader
         train_reader = reader.train()
@@ -132,7 +131,7 @@ def compress(args):
         val_reader, batch_size=args.batch_size, drop_last=True)
     val_program = student_program.clone(for_test=True)
 
-    places = fluid.cuda_places()
+    places = fluid.cuda_places() if args.use_gpu else fluid.cpu_places()
     train_loader.set_sample_list_generator(train_reader, places)
     valid_loader.set_sample_list_generator(val_reader, place)
 
@@ -146,11 +145,13 @@ def compress(args):
                 name='image', shape=image_shape, dtype='float32')
             predict = teacher_model.net(image, class_dim=class_dim)
 
-        #print("="*50+"teacher_model_params"+"="*50)
-        #for v in teacher_program.list_vars():
-        #    print(v.name, v.shape)
+    #print("="*50+"teacher_model_params"+"="*50)
+    #for v in teacher_program.list_vars():
+    #    print(v.name, v.shape)
 
     exe.run(t_startup)
+    _download('http://paddle-imagenet-models-name.bj.bcebos.com/ResNet50_pretrained.tar', '.')
+    _decompress('./ResNet50_pretrained.tar')
     assert args.teacher_pretrained_model and os.path.exists(
         args.teacher_pretrained_model
     ), "teacher_pretrained_model should be set when teacher_model is not None."
@@ -158,7 +159,7 @@ def compress(args):
     def if_exist(var):
         return os.path.exists(
             os.path.join(args.teacher_pretrained_model, var.name)
-        ) and var.name != 'conv1_weights' and var.name != 'fc_0.w_0' and var.name != 'fc_0.b_0'
+        ) and var.name != 'fc_0.w_0' and var.name != 'fc_0.b_0'
 
     fluid.io.load_vars(
         exe,
@@ -173,19 +174,9 @@ def compress(args):
         data_name_map,
         place)
 
-    #print("="*50+"teacher_vars"+"="*50)
-    #for v in teacher_program.list_vars():
-    #    if '_generated_var' not in v.name and 'fetch' not in v.name and 'feed' not in v.name:
-    #        print(v.name, v.shape)
-    #return
-
     with fluid.program_guard(main, s_startup):
-        l2_loss_v = l2_loss("teacher_fc_0.tmp_0", "fc_0.tmp_0", main)
-        fsp_loss_v = fsp_loss("teacher_res2a_branch2a.conv2d.output.1.tmp_0",
-                              "teacher_res3a_branch2a.conv2d.output.1.tmp_0",
-                              "depthwise_conv2d_1.tmp_0", "conv2d_3.tmp_0",
-                              main)
-        loss = avg_cost + l2_loss_v + fsp_loss_v
+        l2_loss = l2_loss("teacher_fc_0.tmp_0", "fc_0.tmp_0", main)
+        loss = avg_cost + l2_loss
         opt = create_optimizer(args)
         opt.minimize(loss)
     exe.run(s_startup)
@@ -196,17 +187,16 @@ def compress(args):
 
     for epoch_id in range(args.num_epochs):
         for step_id, data in enumerate(train_loader):
-            loss_1, loss_2, loss_3, loss_4 = exe.run(
+            loss_1, loss_2, loss_3 = exe.run(
                 parallel_main,
                 feed=data,
                 fetch_list=[
-                    loss.name, avg_cost.name, l2_loss_v.name, fsp_loss_v.name
+                    loss.name, avg_cost.name, l2_loss.name
                 ])
             if step_id % args.log_period == 0:
                 _logger.info(
-                    "train_epoch {} step {} loss {:.6f}, class loss {:.6f}, l2 loss {:.6f}, fsp loss {:.6f}".
-                    format(epoch_id, step_id, loss_1[0], loss_2[0], loss_3[0],
-                           loss_4[0]))
+                    "train_epoch {} step {} loss {:.6f}, class loss {:.6f}, l2 loss {:.6f}".
+                    format(epoch_id, step_id, loss_1[0], loss_2[0], loss_3[0]))
         val_acc1s = []
         val_acc5s = []
         for step_id, data in enumerate(valid_loader):

demo/utility.py

@@ -20,6 +20,12 @@ import distutils.util
 import os
 import numpy as np
 import six
+import requests
+import shutil
+import tqdm
+import hashlib
+import tarfile
+import zipfile
 import logging
 import paddle.fluid as fluid
 import paddle.compat as cpt
@@ -30,6 +36,7 @@ logging.basicConfig(format='%(asctime)s-%(levelname)s: %(message)s')
 _logger = logging.getLogger(__name__)
 _logger.setLevel(logging.INFO)
 
+DOWNLOAD_RETRY_LIMIT=3
 
 def print_arguments(args):
     """Print argparse's arguments.
@@ -154,3 +161,122 @@ def load_persistable_nodes(executor, dirname, graph):
         else:
             _logger.info("Cannot find the var %s!!!" % (node.name()))
     fluid.io.load_vars(executor=executor, dirname=dirname, vars=var_list)
+
+
+def _download(url, path, md5sum=None):
+    """
+    Download from url, save to path.
+
+    url (str): download url
+    path (str): download to given path
+    """
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+    fname = os.path.split(url)[-1]
+    fullname = os.path.join(path, fname)
+    retry_cnt = 0
+
+    while not (os.path.exists(fullname) and _md5check(fullname, md5sum)):
+        if retry_cnt < DOWNLOAD_RETRY_LIMIT:
+            retry_cnt += 1
+        else:
+            raise RuntimeError("Download from {} failed. "
+                               "Retry limit reached".format(url))
+
+        _logger.info("Downloading {} from {}".format(fname, url))
+
+        req = requests.get(url, stream=True)
+        if req.status_code != 200:
+            raise RuntimeError("Downloading from {} failed with code "
+                               "{}!".format(url, req.status_code))
+
+        # For protecting download interupted, download to
+        # tmp_fullname firstly, move tmp_fullname to fullname
+        # after download finished
+        tmp_fullname = fullname + "_tmp"
+        total_size = req.headers.get('content-length')
+        with open(tmp_fullname, 'wb') as f:
+            if total_size:
+                for chunk in tqdm.tqdm(
+                        req.iter_content(chunk_size=1024),
+                        total=(int(total_size) + 1023) // 1024,
+                        unit='KB'):
+                    f.write(chunk)
+            else:
+                for chunk in req.iter_content(chunk_size=1024):
+                    if chunk:
+                        f.write(chunk)
+        shutil.move(tmp_fullname, fullname)
+
+    return fullname
+
+def _md5check(fullname, md5sum=None):
+    if md5sum is None:
+        return True
+
+    _logger.info("File {} md5 checking...".format(fullname))
+    md5 = hashlib.md5()
+    with open(fullname, 'rb') as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            md5.update(chunk)
+    calc_md5sum = md5.hexdigest()
+
+    if calc_md5sum != md5sum:
+        _logger.info("File {} md5 check failed, {}(calc) != "
+                    "{}(base)".format(fullname, calc_md5sum, md5sum))
+        return False
+    return True
+
+def _decompress(fname):
+    """
+    Decompress for zip and tar file
+    """
+    _logger.info("Decompressing {}...".format(fname))
+
+    # For protecting decompressing interupted,
+    # decompress to fpath_tmp directory firstly, if decompress
+    # successed, move decompress files to fpath and delete
+    # fpath_tmp and remove download compress file.
+    fpath = os.path.split(fname)[0]
+    fpath_tmp = os.path.join(fpath, 'tmp')
+    if os.path.isdir(fpath_tmp):
+        shutil.rmtree(fpath_tmp)
+        os.makedirs(fpath_tmp)
+
+    if fname.find('tar') >= 0:
+        with tarfile.open(fname) as tf:
+            tf.extractall(path=fpath_tmp)
+    elif fname.find('zip') >= 0:
+        with zipfile.ZipFile(fname) as zf:
+            zf.extractall(path=fpath_tmp)
+    else:
+        raise TypeError("Unsupport compress file type {}".format(fname))
+
+    for f in os.listdir(fpath_tmp):
+        src_dir = os.path.join(fpath_tmp, f)
+        dst_dir = os.path.join(fpath, f)
+        _move_and_merge_tree(src_dir, dst_dir)
+
+    shutil.rmtree(fpath_tmp)
+    os.remove(fname)
+
+def _move_and_merge_tree(src, dst):
+    """
+    Move src directory to dst, if dst is already exists,
+    merge src to dst
+    """
+    if not os.path.exists(dst):
+        shutil.move(src, dst)
+    else:
+        for fp in os.listdir(src):
+            src_fp = os.path.join(src, fp)
+            dst_fp = os.path.join(dst, fp)
+            if os.path.isdir(src_fp):
+                if os.path.isdir(dst_fp):
+                    _move_and_merge_tree(src_fp, dst_fp)
+                else:
+                    shutil.move(src_fp, dst_fp)
+            elif os.path.isfile(src_fp) and \
+                    not os.path.isfile(dst_fp):
+                shutil.move(src_fp, dst_fp)

doc/demo_guide.md

 
-## [蒸馏]()
+## [蒸馏](../demo/distillation/distillation_demo.py)
+
+蒸馏demo默认使用ResNet50作为teacher网络，MobileNet作为student网络，此外还支持将teacher和student换成[models目录](../demo/models)支持的任意模型。
+
+demo中对teahcer模型和student模型的一层特征图添加了l2_loss的蒸馏损失函数，使用时也可根据需要选择fsp_loss, soft_label_loss以及自定义的loss函数。
+
+训练默认使用的是cifar10数据集，piecewise_decay学习率衰减策略，momentum优化器进行120轮蒸馏训练。使用者也可以简单地用args参数切换为使用ImageNet数据集，cosine_decay学习率衰减策略等其他训练配置。
+