(
self, model,
out_indices=(0, 1, 2, 3, 4), out_map=None, out_as_dict=False, no_rewrite=False,
feature_concat=False, flatten_sequential=False, default_hook_type='forward')
| 246 | FIXME this does not currently work with Torchscript, see FeatureHooks class |
| 247 | """ |
| 248 | def __init__( |
| 249 | self, model, |
| 250 | out_indices=(0, 1, 2, 3, 4), out_map=None, out_as_dict=False, no_rewrite=False, |
| 251 | feature_concat=False, flatten_sequential=False, default_hook_type='forward'): |
| 252 | super(FeatureHookNet, self).__init__() |
| 253 | assert not torch.jit.is_scripting() |
| 254 | self.feature_info = _get_feature_info(model, out_indices) |
| 255 | self.out_as_dict = out_as_dict |
| 256 | layers = OrderedDict() |
| 257 | hooks = [] |
| 258 | if no_rewrite: |
| 259 | assert not flatten_sequential |
| 260 | if hasattr(model, 'reset_classifier'): # make sure classifier is removed? |
| 261 | model.reset_classifier(0) |
| 262 | layers['body'] = model |
| 263 | hooks.extend(self.feature_info.get_dicts()) |
| 264 | else: |
| 265 | modules = _module_list(model, flatten_sequential=flatten_sequential) |
| 266 | remaining = {f['module']: f['hook_type'] if 'hook_type' in f else default_hook_type |
| 267 | for f in self.feature_info.get_dicts()} |
| 268 | for new_name, old_name, module in modules: |
| 269 | layers[new_name] = module |
| 270 | for fn, fm in module.named_modules(prefix=old_name): |
| 271 | if fn in remaining: |
| 272 | hooks.append(dict(module=fn, hook_type=remaining[fn])) |
| 273 | del remaining[fn] |
| 274 | if not remaining: |
| 275 | break |
| 276 | assert not remaining, f'Return layers ({remaining}) are not present in model' |
| 277 | self.update(layers) |
| 278 | self.hooks = FeatureHooks(hooks, model.named_modules(), out_map=out_map) |
| 279 | |
| 280 | def forward(self, x): |
| 281 | for name, module in self.items(): |
nothing calls this directly
no test coverage detected