钩子类,用于捕获前向和后向信息
| 83 | imwrite(output, f"result/cond/test/{step}.png") |
| 84 | |
| 85 | class Hook_back_loop: |
| 86 | """钩子类,用于捕获前向和后向信息""" |
| 87 | def __init__(self, module, module_name): |
| 88 | self.name = module_name |
| 89 | self.hook = module.register_forward_hook(self.forward_hook) |
| 90 | |
| 91 | def forward_hook(self, module, inp, out): |
| 92 | self.input = inp |
| 93 | self.output = out |
| 94 | self.output.register_hook(self.backward_hook) |
| 95 | |
| 96 | def backward_hook(self, grad): |
| 97 | self.grad = grad |
| 98 | return grad |
| 99 | |
| 100 | class BaseAnalysis(BaseModel): |
| 101 | def __init__(self, opt): |