A class used to get weights and activations from specified layers from a Pytorch model.
| 156 | |
| 157 | |
| 158 | class GetWeightAndActivation: |
| 159 | """ |
| 160 | A class used to get weights and activations from specified layers from a Pytorch model. |
| 161 | """ |
| 162 | |
| 163 | def __init__(self, model, layers): |
| 164 | """ |
| 165 | Args: |
| 166 | model (nn.Module): the model containing layers to obtain weights and activations from. |
| 167 | layers (list of strings): a list of layer names to obtain weights and activations from. |
| 168 | Names are hierarchical, separated by /. For example, If a layer follow a path |
| 169 | "s1" ---> "pathway0_stem" ---> "conv", the layer path is "s1/pathway0_stem/conv". |
| 170 | """ |
| 171 | self.model = model |
| 172 | self.hooks = {} |
| 173 | self.layers_names = layers |
| 174 | # eval mode |
| 175 | self.model.eval() |
| 176 | self._register_hooks() |
| 177 | |
| 178 | def _get_layer(self, layer_name): |
| 179 | """ |
| 180 | Return a layer (nn.Module Object) given a hierarchical layer name, separated by /. |
| 181 | Args: |
| 182 | layer_name (str): the name of the layer. |
| 183 | """ |
| 184 | layer_ls = layer_name.split("/") |
| 185 | prev_module = self.model |
| 186 | for layer in layer_ls: |
| 187 | prev_module = prev_module._modules[layer] |
| 188 | |
| 189 | return prev_module |
| 190 | |
| 191 | def _register_single_hook(self, layer_name): |
| 192 | """ |
| 193 | Register hook to a layer, given layer_name, to obtain activations. |
| 194 | Args: |
| 195 | layer_name (str): name of the layer. |
| 196 | """ |
| 197 | |
| 198 | def hook_fn(module, input, output): |
| 199 | self.hooks[layer_name] = output.clone().detach() |
| 200 | |
| 201 | layer = get_layer(self.model, layer_name) |
| 202 | layer.register_forward_hook(hook_fn) |
| 203 | |
| 204 | def _register_hooks(self): |
| 205 | """ |
| 206 | Register hooks to layers in `self.layers_names`. |
| 207 | """ |
| 208 | for layer_name in self.layers_names: |
| 209 | self._register_single_hook(layer_name) |
| 210 | |
| 211 | def get_activations(self, input, bboxes=None): |
| 212 | """ |
| 213 | Obtain all activations from layers that we register hooks for. |
| 214 | Args: |
| 215 | input (tensors, list of tensors): the model input. |