Merge pull request #4 from McCrearyD/accumulated-grad-scoring

Accumulated Gradient Scoring

Author: Dyllan McCreary

rigl_torch/RigL.py

@@ -19,7 +19,16 @@ def __name__(self):
     @torch.no_grad()
     def __call__(self, grad):
         mask = self.scheduler.backward_masks[self.layer]
-        self.dense_grad = grad.clone()
+
+        # only calculate dense_grads when necessary
+        if self.scheduler.check_if_backward_hook_should_accumulate_grad():
+            if self.dense_grad is None:
+                # initialize as all 0s so we can do a rolling average
+                self.dense_grad = torch.zeros_like(grad)
+            self.dense_grad += grad / self.scheduler.grad_accumulation_n
+        else:
+            self.dense_grad = None
+
         return grad * mask
 
 
@@ -34,17 +43,19 @@ def _wrapped_step():
 
 class RigLScheduler:
 
-    def __init__(self, model, optimizer, dense_allocation=1, T_end=None, sparsity_distribution='uniform', ignore_linear_layers=True, is_already_sparsified=False, delta=100, alpha=0.3, static_topo=False):
+    def __init__(self, model, optimizer, dense_allocation=1, T_end=None, sparsity_distribution='uniform', ignore_linear_layers=True, is_already_sparsified=False, delta=100, alpha=0.3, static_topo=False, grad_accumulation_n=1):
         if dense_allocation <= 0 or dense_allocation > 1:
             raise Exception('Dense allocation must be on the interval (0, 1]. Got: %f' % dense_allocation)
 
         self.model = model
         self.optimizer = optimizer
         self.sparsity_distribution = sparsity_distribution
         self.static_topo = static_topo
+        self.grad_accumulation_n = grad_accumulation_n
         self.ignore_linear_layers = ignore_linear_layers
         self.backward_masks = None
 
+        assert self.grad_accumulation_n > 0 and self.grad_accumulation_n < delta
         assert self.sparsity_distribution in ('uniform', )
 
         self.W, self._linear_layers_mask = get_W(model, return_linear_layers_mask=True)
@@ -200,6 +211,19 @@ def apply_mask_to_gradients(self):
 
             w.grad *= mask
 
+    
+    def check_if_backward_hook_should_accumulate_grad(self):
+        """
+        Used by the backward hooks. Basically just checks how far away the next rigl step is, 
+        if it's within `self.grad_accumulation_n` steps, return True.
+        """
+
+        if self.step >= self.T_end:
+            return False
+
+        steps_til_next_rigl_step = self.delta_T - (self.step % self.delta_T)
+        return steps_til_next_rigl_step <= self.grad_accumulation_n
+
 
     def cosine_annealing(self):
         return self.alpha / 2 * (1 + np.cos((self.step * np.pi) / self.T_end))

sagemaker/rigl.ipynb

@@ -11,7 +11,6 @@
     "sagemaker_session = sagemaker.Session()\n",
     "\n",
     "bucket = sagemaker_session.default_bucket()\n",
-    "prefix = 'sagemaker/rigl'\n",
     "\n",
     "role = sagemaker.get_execution_role()"
    ]
@@ -66,6 +65,7 @@
     "                        'static-topo': 0,\n",
     "                        'alpha': 0.3,\n",
     "                        'delta': 100,\n",
+    "#                         'grad-accumulation-n': 4, # if using a smaller batch size, this may be useful\n",
     "                        'batch-size': 1024,\n",
     "                        'lr': 0.1,\n",
     "#                         'lr-warmup-end': 5,\n",
@@ -86,7 +86,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# estimator.fit(file_system_input) # train with FSx Lustre as input\n",
     "estimator.fit('s3://imagenet-compressed-oregon') # use imagenet s3 bucket"
    ]
   }

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="rigl-torch",
-    version="0.2",
+    version="0.3",
     author="Dyllan McCreary",
     author_email="[email protected]",
     description="Implementation of Google Research's \"RigL\" sparse model training method in PyTorch.",

tests/test_rigl.py

@@ -15,9 +15,9 @@
 arch = 'resnet50'
 image_dimensionality = (3, 224, 224)
 num_classes = 1000
-max_iters = 6
+max_iters = 15
 T_end = int(max_iters * 0.75)
-delta = 2
+delta = 3
 dense_allocation = 0.1
 criterion = torch.nn.functional.cross_entropy
 

train_imagenet_rigl.py

@@ -51,6 +51,8 @@
                     help='percentage of dense parameters allowed. if None, pruning will not be used. must be on the interval (0, 1]')
 parser.add_argument('--delta', default=100, type=int,
                     help='delta param for pruning')
+parser.add_argument('--grad-accumulation-n', default=1, type=int,
+                    help='number of gradients to accumulate before scoring for rigl')
 parser.add_argument('--alpha', default=0.3, type=float,
                     help='alpha param for pruning')
 parser.add_argument('--static-topo', default=0, type=int, help='if 1, use random sparsity topo and remain static')
@@ -291,7 +293,7 @@ def main_worker(gpu, ngpus_per_node, args):
     if args.dense_allocation is not None:
         total_iterations = args.epochs * len(train_loader)
         T_end = int(0.75 * total_iterations) # (stop tweaking topology after 75% of training)
-        pruner = RigLScheduler(model, optimizer, dense_allocation=args.dense_allocation, T_end=T_end, delta=args.delta, alpha=args.alpha, static_topo=args.static_topo)
+        pruner = RigLScheduler(model, optimizer, dense_allocation=args.dense_allocation, T_end=T_end, delta=args.delta, alpha=args.alpha, static_topo=args.static_topo, grad_accumulation_n=args.grad_accumulation_n)
         print('pruning with dense allocation: %f & T_end=%i' % (args.dense_allocation, T_end))
         print(pruner)
 

train_mnist_rigl.py

@@ -91,6 +91,8 @@ def main():
                         help='percentage of dense parameters allowed. if None, pruning will not be used. must be on the interval (0, 1]')
     parser.add_argument('--delta', default=100, type=int,
                         help='delta param for pruning')
+    parser.add_argument('--grad-accumulation-n', default=1, type=int,
+                        help='number of gradients to accumulate before scoring for rigl')
     parser.add_argument('--alpha', default=0.3, type=float,
                         help='alpha param for pruning')
     parser.add_argument('--static-topo', default=0, type=int, help='if 1, use random sparsity topo and remain static')
@@ -154,7 +156,7 @@ def main():
     pruner = lambda: True
     if args.dense_allocation is not None:
         T_end = int(0.75 * args.epochs * len(train_loader))
-        pruner = RigLScheduler(model, optimizer, dense_allocation=args.dense_allocation, alpha=args.alpha, delta=args.delta, static_topo=args.static_topo, T_end=T_end, ignore_linear_layers=False)
+        pruner = RigLScheduler(model, optimizer, dense_allocation=args.dense_allocation, alpha=args.alpha, delta=args.delta, static_topo=args.static_topo, T_end=T_end, ignore_linear_layers=False, grad_accumulation_n=args.grad_accumulation_n)
 
     print(model)
     for epoch in range(1, args.epochs + 1):