Can hook remove itself?

blackbirdbarber: <p>After it runs can hook remove itself from inside?</p>

ptrblck: <p>If you register a hook, you should get a handle to this hook, which can be used to remove it.</p>

blackbirdbarber: <p>Yes, this handle is OK, but what I actually wanted to check if it is possible to do it from the function itself. Some easy and smart way.</p>
<p>Because usually this is what I need for the hook to run once.</p>

ptrblck: <p>You could try to access the <code>handle</code> from inside the hook:</p>
<pre><code class="lang-python">def hook(grad):
    print(grad)
    handle.remove()
    
model = nn.Linear(1, 1)
handle = model.weight.register_hook(hook)

model(torch.randn(1, 1)).backward()
model(torch.randn(1, 1)).backward()
</code></pre>
<p>This code will print the gradient only once and remove the hook.</p>

111577: <p>How can I remove the hooks if they are like this?</p>
<pre><code class="lang-auto">def register_hooks(self):
        """Register forward and backward hook to Conv module."""
        for module, name in self.conv_names.items():
            module.register_forward_hook(self.save_input_forward_hook)
            module.register_full_backward_hook(self.compute_fisher_backward_hook)

self.register_hooks()
</code></pre>

ptrblck: <p>You could store the <code>handle</code>, which is returned by the <code>register..._hook</code> methods, and remove it when needed.</p>

111577: <p>I did the following</p>
<pre><code class="lang-auto">def register_hooks(self):
"""Register forward and backward hook to Conv module."""
    self.forward_hooks, self.backward_hooks= [], []
    for module, name in self.conv_names.items():
        self.forward_hooks.append(module.register_forward_hook(self.save_input_forward_hook))
        self.backward_hooks.append(module.register_full_backward_hook(self.compute_fisher_backward_hook)

def remove_hooks(self):
    for fh in self.forward_hooks:
        fh.remove()
    for bh in self.backward_hooks:
        bh.remove() 
</code></pre>
<p>Then I remove hooks and try to save the model</p>
<pre><code class="lang-auto">torch.save(deepcopy(model).half(), path)
</code></pre>
<p>and I have the following error</p>
<pre><code class="lang-auto">AttributeError: 'Conv2d' object has no attribute 'modified_forward'
</code></pre>
<p>I will be grateful if you help me! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>Could you post an executable code snippet, which would reproduce the issue, please?</p>

111577: <aside class="quote no-group" data-username="ptrblck" data-post="8" data-topic="55266">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>executable code snippet</p>
</blockquote>
</aside>
<p>Sorry, I was wrong. I can save the model. But then I try to load this model. And I have this error.<br>
There is code snippet, where model is loading.</p>
<pre><code class="lang-auto">    if RANK in [-1, 0]:
        LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
        for f in last, best:
            if f.exists():
                strip_optimizer(f)  # strip optimizers
                if f is best:
                    LOGGER.info(f'\nValidating {f}...')
                    results, _, _ = val.run(data_dict,
                                            batch_size=batch_size // WORLD_SIZE * 2,
                                            imgsz=imgsz,
                                            model=attempt_load(f, device).half(),
                                            iou_thres=0.65 if is_coco else 0.60,  # best pycocotools results at 0.65
                                            single_cls=single_cls,
                                            dataloader=val_loader,
                                            save_dir=save_dir,
                                            save_json=is_coco,
                                            verbose=True,
                                            plots=True,
                                            callbacks=callbacks,
                                            compute_loss=compute_loss)  
</code></pre>
<pre><code class="lang-auto">def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_optimizer()
    # Strip optimizer from 'f' to finalize training, optionally save as 's'
    x = torch.load(f, map_location=torch.device('cpu'))
    if x.get('ema'):
        x['model'] = x['ema']  # replace model with ema
    for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates':  # keys
        x[k] = None
    x['epoch'] = -1
    x['model'].half()  # to FP16
    for p in x['model'].parameters():
        p.requires_grad = False
    torch.save(x, s or f)
    mb = os.path.getsize(s or f) / 1E6  # filesize
    print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
</code></pre>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "train.py", line 651, in &lt;module&gt;
    main(opt)
  File "train.py", line 548, in main
    train(opt.hyp, opt, device, callbacks)
  File "train.py", line 437, in train
    strip_optimizer(f)  # strip optimizers
  File "/home/apleshkova/yolov5/utils/general.py", line 703, in strip_optimizer
    x = torch.load(f, map_location=torch.device('cpu'))
  File "/home/apleshkova/.local/lib/python3.6/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/apleshkova/.local/lib/python3.6/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
  File "/home/apleshkova/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1131, in __getattr__
    type(self).__name__, name))
AttributeError: 'Conv2d' object has no attribute 'modified_forward'

</code></pre>

ptrblck: <p>Unfortunately, the code is not executable so I can’t debug it but can give my best guesses.<br>
It seems you are storing the entire model instead of the <code>state_dict</code>, which I would not recommend as I’ve seen it failing in various ways.<br>
Store the <code>state_dict</code> instead, recreate the model instance, and load the <code>state_dict</code> back afterwards.</p>

111577: <p>Thank you for your advice! This solved my problem)</p>