Register forward pass hook (which registers a backward hook) to model outputs Returns: - layers: a dict with keys as layer/module and values as layer/module names e.g. layers[nn.Conv2d] = layer1.0.conv1 - grads: a list of tuples with module name and tensor outp
(model, hook_forward, hook_backward)
| 154 | return hook |
| 155 | |
| 156 | def get_all_layers(model, hook_forward, hook_backward): |
| 157 | """Register forward pass hook (which registers a backward hook) to model outputs |
| 158 | |
| 159 | Returns: |
| 160 | - layers: a dict with keys as layer/module and values as layer/module names |
| 161 | e.g. layers[nn.Conv2d] = layer1.0.conv1 |
| 162 | - grads: a list of tuples with module name and tensor output gradient |
| 163 | e.g. grads[0] == (layer1.0.conv1, tensor.Torch(...)) |
| 164 | """ |
| 165 | layers = dict() |
| 166 | grads = [] |
| 167 | for name, layer in model.named_modules(): |
| 168 | # skip Sequential and/or wrapper modules |
| 169 | if any(layer.children()) is False: |
| 170 | layers[layer] = name |
| 171 | layer.register_forward_hook(hook_forward(name, grads, hook_backward)) |
| 172 | return layers, grads |
| 173 | |
| 174 | # register hooks |
| 175 | layers_bn, grads_bn = get_all_layers(model_bn, hook_forward, hook_backward) |
no test coverage detected