(self, module, input, output)
| 212 | self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n") |
| 213 | |
| 214 | def create_frame(self, module, input, output): |
| 215 | self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}") |
| 216 | |
| 217 | # params |
| 218 | for name, p in module.named_parameters(recurse=False): |
| 219 | self.analyse_variable(p, name) |
| 220 | |
| 221 | # inputs |
| 222 | if isinstance(input, tuple): |
| 223 | for i, x in enumerate(input): |
| 224 | self.analyse_variable(x, f"input[{i}]") |
| 225 | else: |
| 226 | self.analyse_variable(input, "input") |
| 227 | |
| 228 | # outputs |
| 229 | if isinstance(output, tuple): |
| 230 | for i, x in enumerate(output): |
| 231 | # possibly a tuple of tuples |
| 232 | if isinstance(x, tuple): |
| 233 | for j, y in enumerate(x): |
| 234 | self.analyse_variable(y, f"output[{i}][{j}]") |
| 235 | else: |
| 236 | self.analyse_variable(x, f"output[{i}]") |
| 237 | else: |
| 238 | self.analyse_variable(output, "output") |
| 239 | |
| 240 | self.save_frame() |
| 241 | |
| 242 | def register_forward_hook(self): |
| 243 | self.model.apply(self._register_forward_hook) |
no test coverage detected