Closed jlmckins closed 2 years ago
Here is the complete error message:
Using pytorch model zoo pretrained model as teacher: resnet18 Traceback (most recent call last): File "train_resnet_imagenet_stepsizex4init_hiway4.py", line 1317, in <module> calc_hutchison_trace(teacher, simplecriterion) File "train_resnet_imagenet_stepsizex4init_hiway4.py", line 1261, in calc_hutchison_trace loss.backward() File "/u/jlmckins/.local/lib/python3.7/site-packages/torch/_tensor.py", line 307, in backward torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs) File "/u/jlmckins/.local/lib/python3.7/site-packages/torch/autograd/__init__.py", line 156, in backward allow_unreachable=True, accumulate_grad=True) # allow_unreachable flag File "/u/jlmckins/.local/lib/python3.7/site-packages/torch/utils/hooks.py", line 110, in hook res = user_hook(self.module, grad_input, self.grad_outputs) File "/u/jlmckins/.local/lib/python3.7/site-packages/backpack/__init__.py", line 209, in hook_run_extensions backpack_extension(module, g_inp, g_out) File "/u/jlmckins/.local/lib/python3.7/site-packages/backpack/extensions/backprop_extension.py", line 125, in __call__ module_extension = self.__get_module_extension(module) File "/u/jlmckins/.local/lib/python3.7/site-packages/backpack/extensions/backprop_extension.py", line 100, in __get_module_extension f"Extension saving to {self.savefield} " NotImplementedError: Extension saving to diag_h does not have an extension for Module <class 'backpack.custom_module.branching.SumModule'>
Hi @jlmckins,
BackPACK's DiagHessian
extension does not support branched computation graphs like ResNets yet. That is because computing the Hessian diagonal in the presence of arbitrary branching becomes a headache.
For your specific case, there's a way out: Since you're using a ResNet with identity skip connections and ReLUs, the Hessian trace is identical to the GGN trace. You can thus use BackPACK's DiagGGNExact
that supports ResNets, and sum the diagonal to obtain the trace.
Best, Felix
Thank you for the fast response! I will give it a try.
I was able to run the 2nd order examples after using the. extend(model, use_converter=True) for ResNet18. However, when I try to run the Hutchinson trace example, I get the following error:
NotImplementedError: Extension saving to diag_h does not have an extension for Module <class 'backpack.custom_module.branching.SumModule'>
Is it possible to extend this module in order to be able to compute the Hutchison trace layerwise for ResNet models?
Thank you, Jeff
Here is part of the test code:
`def calc_hutchison_trace(model, criterion):
`