fix

Author: xiaosu-zhu

configs/hash/CSQ_D/easy/cifar10/32-bits.yaml

@@ -45,7 +45,7 @@ train:
   schdr:
     key: CosineAnnealingLR
     params:
-      T_max: 100
+      T_max: 500
   gpu:
     gpus: 1
     vRam: -1

configs/quantization/PQNet/easy/cifar10/32-bits.yaml

@@ -43,8 +43,7 @@ train:
   schdr:
     key: CosineAnnealingLR
     params:
-      T_max: 100
-      eta_min: 1.e-8
+      T_max: 500
   gpu:
     gpus: 1
     vRam: -1

modfire/criterion/csq.py

@@ -102,6 +102,12 @@ def __init__(self, bits) -> None:
             self._net = nn.ModuleList(ffnNet() for _ in range(bits // 8))
             self._bitFlip = CSQ_D._randomBitFlip(bits, int(bits // 32) ** 2)
 
+        def reset(self):
+            for ffnNet in self._net:
+                for net in ffnNet:
+                    if hasattr(net, "reset_parameters"):
+                        net.reset_parameters()
+
         def forward(self, x, flip):
             if flip:
                 x = self._bitFlip(x)
@@ -132,14 +138,16 @@ def BitFlip(self, numBitsToFlip: int):
         self.mapper._bitFlip.BitFlip = numBitsToFlip
         self.bitFlip.BitFlip = numBitsToFlip
 
-    def resetPermIdx(self):
+    def reset(self):
         # reset permIdx
         self.permIdx.data.copy_(torch.randperm(self.m * 8, device=self.permIdx.device))
+        # reset params
+        self.mapper.reset()
 
     def forward(self, x: torch.Tensor, y: torch.Tensor):
         self._ticker += 1
         if self._ticker % int(self.m * 16) == 0:
-            self.resetPermIdx()
+            self.reset()
         # X are permuted on last dim according to permIdx
         x = x[:, self.permIdx]