| 77 | |
| 78 | # step 3: catch the input of thefirst layer |
| 79 | class Catcher(nn.Module): |
| 80 | def __init__(self, module, dataset): |
| 81 | super().__init__() |
| 82 | self.module = module |
| 83 | self.dataset = dataset |
| 84 | self.index = 0 |
| 85 | self.attention_mask = None |
| 86 | self.position_ids = None |
| 87 | |
| 88 | def forward(self, inp, **kwargs): |
| 89 | self.dataset.update_data(self.index, inp.squeeze(0).to('cpu')) |
| 90 | self.index += 1 |
| 91 | if self.attention_mask is None: |
| 92 | self.attention_mask = kwargs["attention_mask"] |
| 93 | if self.position_ids is None: |
| 94 | self.position_ids = kwargs["position_ids"] |
| 95 | raise ValueError |
| 96 | |
| 97 | # step 3.1: catch the input of training set |
| 98 | layers[0] = Catcher(layers[0],fp_train_inps) |