def backward_hook(module: nn.Module, grad_input: torch.Tensor, grad_output: torch.Tensor) -> None: # print("Backward Module : {}, grad_inp: {}, grad_out: {}".format(module, len(grad_input), len(grad_output))) # Forward pass: compute predicted y using operations; we compute # ReLU using our custom autograd operation. if isinstance(module, nn.Conv2d): inp = QuantizationCrunch[str(module.__hash__())]["input"] # module.register_buffer("orig_output", output) # torch.utils.hooks.RemovableHandle(clone).remove() # print("w: {}, b: {}".format(module.weight, module.bias)) quant_out, quant_weights = quantize(module, module.weight, module.bias, inp)
def forward_hook(module: nn.Module, input: tuple, output: torch.Tensor) -> None: # print("Forward Module : {}. hash {}".format(module, module.__hash__())) # for layer in module.modules(): if isinstance(module, nn.Conv2d): # print("Layer dict {}".format(module.state_dict().keys())) # str = "quant_{}_input".format(list(module.named_modules()) # module.register_parameter("orig_input", input[0]) # module.register_buffer("orig_output", output[0]) QuantizationCrunch[str(module.__hash__())] = { "input": input[0], "output": output[0] } # module.