| 27 | self.scaler_row = torch.zeros((self.columns), device=self.dev) |
| 28 | |
| 29 | def add_batch(self, inp, out): |
| 30 | if len(inp.shape) == 2: |
| 31 | inp = inp.unsqueeze(0) |
| 32 | tmp = inp.shape[0] |
| 33 | if isinstance(self.layer, nn.Linear) or isinstance( |
| 34 | self.layer, transformers.Conv1D |
| 35 | ): |
| 36 | if len(inp.shape) == 3: |
| 37 | inp = inp.reshape((-1, inp.shape[-1])) |
| 38 | inp = inp.t() |
| 39 | self.H *= self.nsamples / (self.nsamples + tmp) |
| 40 | |
| 41 | self.scaler_row *= self.nsamples / (self.nsamples + tmp) |
| 42 | |
| 43 | self.nsamples += tmp |
| 44 | inp = math.sqrt(2 / self.nsamples) * inp.float() |
| 45 | self.H += inp.matmul(inp.t()) |
| 46 | self.scaler_row += torch.norm(inp, p=2, dim=1) ** 2 / self.nsamples |
| 47 | |
| 48 | def get_wanda_mask(self, sparsity, prunen, prunem): |
| 49 | W_metric = torch.abs(self.layer.weight.data) * torch.sqrt( |