1.0rc

Author: GT-ZhangAcer

dev.py

@@ -16,11 +16,6 @@
 # 导入RIFLE模块
 from paddle_rifle.rifle import RIFLECallback
 
-transform = Compose([ToTensor()])
-
-train_data = MNIST(transform=transform)
-test_data = MNIST(mode="test", transform=transform)
-
 
 class Net(nn.Layer):
     def __init__(self, num_classes=10):
@@ -43,24 +38,37 @@ def forward(self, inputs):
         return x
 
 
-net = Net(num_classes=10)
-fc_layer = net.fc2
+def main(use_init: bool = False):
+    transform = Compose([ToTensor()])
+
+    train_data = MNIST(transform=transform)
+    test_data = MNIST(mode="test", transform=transform)
+
+    net = Net(num_classes=10)
+    fc_layer = net.fc2
+
+    model = paddle.Model(network=net,
+                         inputs=paddle.static.InputSpec([1, 28, 28], name="ipt"),
+                         labels=paddle.static.InputSpec([1], dtype="int64", name="lab"))
+
+    rifle_cb = RIFLECallback(fc_layer,
+                             re_init_epoch=1,
+                             max_re_num=3,
+                             weight_initializer=paddle.nn.initializer.XavierNormal() if use_init else None)
 
-model = paddle.Model(network=net,
-                     inputs=paddle.static.InputSpec([1, 28, 28], name="ipt"),
-                     labels=paddle.static.InputSpec([1], dtype="int64", name="lab"))
+    sgd = paddle.optimizer.SGD(parameters=model.parameters())
+    loss = paddle.nn.loss.CrossEntropyLoss()
+    acc = paddle.metric.Accuracy((1, 5))
+    model.prepare(sgd, loss, acc)
 
-rifle_cb = RIFLECallback(fc_layer, 1, 3, weight_initializer=paddle.nn.initializer.XavierNormal())
+    # 开始训练并传入RIFLE Callback
+    model.fit(train_data,
+              test_data,
+              batch_size=256,
+              epochs=2,
+              log_freq=10,
+              callbacks=[rifle_cb])
 
-sgd = paddle.optimizer.SGD(parameters=model.parameters())
-loss = paddle.nn.loss.CrossEntropyLoss()
-acc = paddle.metric.Accuracy((1, 5))
-model.prepare(sgd, loss, acc)
 
-# 开始训练并传入RIFLE Callback
-model.fit(train_data,
-          test_data,
-          batch_size=32,
-          epochs=20,
-          log_freq=100,
-          callbacks=[rifle_cb])
+if __name__ == '__main__':
+    main()

paddle_rifle/rifle.py

@@ -2,7 +2,6 @@
 # Datetime: 2021/2/23 
 # Copyright belongs to the author.
 # Please indicate the source for reprinting.
-import copy
 
 import paddle
 
@@ -40,6 +39,7 @@ def __init__(self,
                  weight_initializer=None):
         """
         :param layers: 需要重置的Layer 或 Layer列表
+        :type layers: (paddle.nn.Layer|list)
         :param re_init_epoch: 经历多少EPOCH后重新初始化输出层
         :param max_re_num: Layer最大重置次数
         :param weight_initializer: 权重默认初始化方案（若为None则为原始权重，可为paddle.nn.initializer.XavierNormal()）
@@ -58,41 +58,27 @@ def apply(self, current_epoch: int):
         应用RIFLE
         :param current_epoch: 当前遍历过的EPOCH数量
         """
-        if current_epoch % self.re_init_epoch == 0 and (current_epoch // self.re_init_epoch) <= self.max_re_num:
-            print_str = f"Initialization successful, {len(self.layers)} group layers will apply RIFLE"
-
-            if current_epoch == 0 and self.weight_initializer is None:
-                for layer_id, layer in enumerate(self.layers):
-                    for param_id, param in enumerate(layer.parameters()):
-                        if ".w_" in param.name:
-                            self.CACHE_PARAMS[f"w_{layer_id}_{param_id}"] = layer.weight.numpy()
-                        elif ".b_" in param.name:
-                            self.CACHE_PARAMS[f"b_{layer_id}_{param_id}"] = layer.bias.numpy()
-            else:
-                for layer_id, layer in enumerate(self.layers):
+        if current_epoch == 0:
+            print(f"\033[0;37;41mInitialization successful, {len(self.layers)} group layers will apply RIFLE\033[0m")
+        elif current_epoch % self.re_init_epoch == 0 and (current_epoch // self.re_init_epoch) <= self.max_re_num:
+            for layer_id, layer in enumerate(self.layers):
+                if self.weight_initializer is None:
+                    layer.parameters().clear()
+                else:
                     for param_id, param in enumerate(layer.parameters()):
                         if ".w_" in param.name:
-                            if self.weight_initializer is not None:
-                                layer.weight = layer.create_parameter(shape=layer.weight.shape,
-                                                                      attr=None,
-                                                                      dtype=layer.weight.dtype,
-                                                                      is_bias=False,
-                                                                      default_initializer=self.weight_initializer)
-                            else:
-                                layer.weight.set_value(self.CACHE_PARAMS[f"w_{layer_id}_{param_id}"])
-
+                            layer.weight = layer.create_parameter(shape=layer.weight.shape,
+                                                                  attr=None,
+                                                                  dtype=layer.weight.dtype,
+                                                                  is_bias=False,
+                                                                  default_initializer=self.weight_initializer)
                         elif ".b_" in param.name:
-                            if self.weight_initializer is not None:
-                                layer.bias = layer.create_parameter(shape=layer.bias.shape,
-                                                                    attr=None,
-                                                                    dtype=layer.bias.dtype,
-                                                                    is_bias=True)
-                            else:
-                                layer.bias.set_value(self.CACHE_PARAMS[f"b_{layer_id}_{param_id}"])
-
-                print_str = f"RIFLE:  layer has been reset in the {current_epoch} epoch!"
+                            layer.bias = layer.create_parameter(shape=layer.bias.shape,
+                                                                attr=None,
+                                                                dtype=layer.bias.dtype,
+                                                                is_bias=True)
 
-            print(f"\033[0;37;41m{print_str}\033[0m")
+            print(f"\033[0;37;41mRIFLE:  layer has been reset in the {current_epoch} epoch!\033[0m")
 
 
 class RIFLECallback(paddle.callbacks.Callback):
@@ -126,6 +112,7 @@ def __init__(self,
         """
         RIFLE的CallBack实现
         :param layers: 需要进行RIFLE的输出层
+        :type layers: (paddle.nn.Layer|list)
         :param re_init_epoch: 经历多少EPOCH后重新初始化输出层
         :param max_re_num: Layer最大重置次数
         :param weight_initializer: 权重默认初始化方案（若为None则为原始权重，可为paddle.nn.initializer.XavierNormal()）

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='paddle-rifle',
-    version='0.21',
+    version='1.0rc',
     packages=["paddle_rifle"],
     url='https://github.com/GT-ZhangAcer/RIFLE_Module',
     license='MIT',

unit_test.py

@@ -0,0 +1,14 @@
+import unittest
+from dev import main
+
+
+class MyTestCase(unittest.TestCase):
+    def test_none_init(self):
+        main()
+
+    def test_init(self):
+        main(True)
+
+
+if __name__ == '__main__':
+    unittest.main()