Feat (graph/gpfq): compression with random projection (#964)

src/brevitas/graph/gpfq.py

@@ -2,10 +2,14 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from copy import deepcopy
+import math
+from math import pi
 from typing import Callable, List, Optional
 
 import numpy as np
 import torch
+from torch.fft import fft
+from torch.fft import fftn
 import torch.nn as nn
 import unfoldNd
 
@@ -19,6 +23,24 @@
 import brevitas.nn as qnn
 
 
+def random_projection(
+ float_input: torch.Tensor, quantized_input: torch.Tensor, compression_rate: float):
+ # use random projection to reduce dimensionality
+ n = quantized_input.size(1)
+ target_dim = int(compression_rate * n)
+ dev = float_input.device
+ # create gaussian random matrix
+ R = torch.normal(mean=0.0, std=1. / math.sqrt(n), size=(target_dim, n), device=dev)
+ quantized_input = torch.transpose(quantized_input, 1, 2) @ R.T
+ float_input = torch.transpose(float_input, 1, 2) @ R.T
+ del R
+ # reshape back
+ quantized_input = torch.transpose(quantized_input, 1, 2)
+ float_input = torch.transpose(float_input, 1, 2)
+
+ return float_input, quantized_input
+
+
 class gpfq_mode(gpxq_mode):
  """
  Apply GPFQ algorithm.
@@ -64,7 +86,8 @@ def __init__(
  act_order: bool = False,
  use_gpfa2q: bool = False,
  accumulator_bit_width: Optional[int] = None,
- a2q_layer_filter_fnc: Optional[Callable[[nn.Module], bool]] = lambda x: True) -> None:
+ a2q_layer_filter_fnc: Optional[Callable[[nn.Module], bool]] = lambda x: True,
+ compression_rate: Optional[float] = 0.0) -> None:
  if not inplace:
  model = deepcopy(model)
  super().__init__(
@@ -83,6 +106,11 @@ def __init__(
  self.accumulator_bit_width = accumulator_bit_width
  self.a2q_layer_filter_fnc = a2q_layer_filter_fnc # returns true when to use GPFA2Q
 
+ # selecting impl of random proj
+ self.compression_rate = compression_rate
+ if self.compression_rate < 0.0 or self.compression_rate > 1.0:
+ raise ValueError('Compression rate for random projection must be between 0 and 1.')
+
  def catch_stopfwd(self, *args, **kwargs):
  # Collect quant input
  try:
@@ -127,7 +155,8 @@ def initialize_module_optimizer(
  act_order=act_order,
  len_parallel_layers=len_parallel_layers,
  create_weight_orig=create_weight_orig,
- p=self.p)
+ p=self.p,
+ compression_rate=self.compression_rate)
  else:
  return GPFA2Q(
  layer=layer,
@@ -136,22 +165,26 @@ def initialize_module_optimizer(
  len_parallel_layers=len_parallel_layers,
  create_weight_orig=create_weight_orig,
  p=self.p,
- accumulator_bit_width=self.accumulator_bit_width)
+ accumulator_bit_width=self.accumulator_bit_width,
+ compression_rate=self.compression_rate)
 
 
 class GPFQ(GPxQ):
  """
  Based on https://github.com/YixuanSeanZhou/Quantized_Neural_Nets/tree/main
  """
 
- def __init__(self, layer, name, act_order, len_parallel_layers, create_weight_orig, p) -> None:
+ def __init__(
+ self, layer, name, act_order, len_parallel_layers, create_weight_orig, p,
+ compression_rate) -> None:
 
  super().__init__(layer, name, act_order, len_parallel_layers, create_weight_orig)
 
  self.float_input = None
  self.quantized_input = None
  self.index_computed = False
  self.p = p
+ self.compression_rate = compression_rate
 
  def update_batch(self, module, input, current_layer):
  if self.disable_pre_forward_hook:
@@ -246,10 +279,12 @@ def single_layer_update(self):
  weight = weight.transpose(1, 0) # This performs a view
  weight = weight.flatten(1)
  weight = weight.view(self.groups, -1, weight.shape[-1]) # [Groups, OC/Groups, IC]
- U = torch.zeros(
- weight.shape[0], weight.shape[1], self.float_input.shape[1], device=dev, dtype=dtype)
+ if self.compression_rate > 0.0:
+ self.float_input, self.quantized_input = random_projection(self.float_input, self.quantized_input, self.compression_rate)
  self.float_input = self.float_input.to(dev)
  self.quantized_input = self.quantized_input.to(dev)
+ U = torch.zeros(
+ weight.shape[0], weight.shape[1], self.float_input.shape[1], device=dev, dtype=dtype)
  # We don't need full Hessian, we just need the diagonal
  self.H_diag = self.quantized_input.transpose(2, 1).square().sum(
  2) # summing over Batch dimension
@@ -300,15 +335,17 @@ def __init__(
  len_parallel_layers,
  create_weight_orig,
  accumulator_bit_width,
- p) -> None:
+ p,
+ compression_rate) -> None:
  GPFQ.__init__(
  self,
  layer=layer,
  name=name,
  act_order=act_order,
  len_parallel_layers=len_parallel_layers,
  create_weight_orig=create_weight_orig,
- p=p)
+ p=p,
+ compression_rate=compression_rate)
  self.accumulator_bit_width = accumulator_bit_width
  assert self.accumulator_bit_width is not None
 
@@ -329,6 +366,8 @@ def single_layer_update(self):
  weight = weight.view(self.groups, -1, weight.shape[-1]) # [Groups, OC/Groups, IC]
  U = torch.zeros(
  weight.shape[0], weight.shape[1], self.float_input.shape[1], device=dev, dtype=dtype)
+ if self.compression_rate > 0.0:
+ self.float_input, self.quantized_input = random_projection(self.float_input, self.quantized_input, self.compression_rate)
  self.float_input = self.float_input.to(dev)
  self.quantized_input = self.quantized_input.to(dev)
 

src/brevitas_examples/imagenet_classification/ptq/ptq_common.py

@@ -535,7 +535,14 @@ def apply_gptq(calib_loader, model, act_order=False):
  gptq.update()
 
 
-def apply_gpfq(calib_loader, model, act_order, p=1.0, use_gpfa2q=False, accumulator_bit_width=None):
+def apply_gpfq(
+ calib_loader,
+ model,
+ act_order,
+ p=1.0,
+ use_gpfa2q=False,
+ accumulator_bit_width=None,
+ compression_rate=0.0):
  model.eval()
  dtype = next(model.parameters()).dtype
  device = next(model.parameters()).device
@@ -545,7 +552,8 @@ def apply_gpfq(calib_loader, model, act_order, p=1.0, use_gpfa2q=False, accumula
  use_quant_activations=True,
  act_order=act_order,
  use_gpfa2q=use_gpfa2q,
- accumulator_bit_width=accumulator_bit_width) as gpfq:
+ accumulator_bit_width=accumulator_bit_width,
+ compression_rate=compression_rate) as gpfq:
  gpfq_model = gpfq.model
  for i in tqdm(range(gpfq.num_layers)):
  for i, (images, target) in enumerate(calib_loader):

src/brevitas_examples/imagenet_classification/ptq/ptq_evaluate.py

@@ -239,6 +239,12 @@ def parse_type(v, default_type):
  help=
  'Split Ratio for Channel Splitting. When set to 0.0, Channel Splitting will not be applied. (default: 0.0)'
 )
+parser.add_argument(
+ '--compression-rate',
+ default=0.0,
+ type=float,
+ help='Specify compression rate < 1.0 for random projection. Default is 0.0 and does not use RP.'
+)
 add_bool_arg(parser, 'gptq', default=False, help='GPTQ (default: disabled)')
 add_bool_arg(parser, 'gpfq', default=False, help='GPFQ (default: disabled)')
 add_bool_arg(parser, 'gpfa2q', default=False, help='GPFA2Q (default: disabled)')
@@ -426,7 +432,12 @@ def main():
 
  if args.gpfq:
  print("Performing GPFQ:")
- apply_gpfq(calib_loader, quant_model, p=args.gpfq_p, act_order=args.gpxq_act_order)
+ apply_gpfq(
+ calib_loader,
+ quant_model,
+ p=args.gpfq_p,
+ act_order=args.gpxq_act_order,
+ compression_rate=args.compression_rate)
 
  if args.gpfa2q:
  print("Performing GPFA2Q:")
@@ -436,7 +447,8 @@ def main():
  p=args.gpfq_p,
  act_order=args.gpxq_act_order,
  use_gpfa2q=args.gpfa2q,
- accumulator_bit_width=args.accumulator_bit_width)
+ accumulator_bit_width=args.accumulator_bit_width,
+ compression_rate=args.compression_rate)
 
  if args.gptq:
  print("Performing GPTQ:")