(
self, model,
out_indices=(0, 1, 2, 3, 4), out_map=None, feature_concat=False, flatten_sequential=False)
| 175 | flatten_sequential (bool): whether to flatten sequential modules assigned to model |
| 176 | """ |
| 177 | def __init__( |
| 178 | self, model, |
| 179 | out_indices=(0, 1, 2, 3, 4), out_map=None, feature_concat=False, flatten_sequential=False): |
| 180 | super(FeatureDictNet, self).__init__() |
| 181 | self.feature_info = _get_feature_info(model, out_indices) |
| 182 | self.concat = feature_concat |
| 183 | self.return_layers = {} |
| 184 | return_layers = _get_return_layers(self.feature_info, out_map) |
| 185 | modules = _module_list(model, flatten_sequential=flatten_sequential) |
| 186 | remaining = set(return_layers.keys()) |
| 187 | layers = OrderedDict() |
| 188 | for new_name, old_name, module in modules: |
| 189 | layers[new_name] = module |
| 190 | if old_name in remaining: |
| 191 | # return id has to be consistently str type for torchscript |
| 192 | self.return_layers[new_name] = str(return_layers[old_name]) |
| 193 | remaining.remove(old_name) |
| 194 | if not remaining: |
| 195 | break |
| 196 | assert not remaining and len(self.return_layers) == len(return_layers), \ |
| 197 | f'Return layers ({remaining}) are not present in model' |
| 198 | self.update(layers) |
| 199 | |
| 200 | def _collect(self, x) -> (Dict[str, torch.Tensor]): |
| 201 | out = OrderedDict() |
nothing calls this directly
no test coverage detected