Repository: wandb/wandb

Issue:
[CLI] `wandb.watch` causes `RuntimeError: Output 0 of BackwardHookFunctionBackward is a view and is being modified inplace.` during the forward pass
**Description**
Calling `wandb.watch` on a PyTorch model causes a `RuntimeError` during the forward pass of said model. The error seems to be related to hooks:
```
RuntimeError: Output 0 of BackwardHookFunctionBackward is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can remove this warning by cloning the output of the custom Function.
```

**Wandb features**
`wandb.watch` on a PyTorch model

**How to reproduce**
Running the code below reproduces the error:
```python
import wandb
from torch.utils.data import DataLoader
from torchvision.models import resnet18
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor


def main():
    data = FakeData(transform=ToTensor())
    dataloader = DataLoader(data)
    model = resnet18()
    wandb.init()
    wandb.watch(model)
    input, _ = next(iter(dataloader))
    model(input)


if __name__ == '__main__':
    main()
```
By omitting the `wandb.watch(model)` line, there are no errors and the code executes successfully.

**Environment**
- OS: Ubuntu 18.04.5 LTS (x86_64)
- Environment: conda
- Python version: 3.8 (64-bit runtime)
- wandb version: 0.10.29
- PyTorch version: 1.9.0.dev20210502
- Torchaudio version: 0.9.0.dev20210502
- Torchvision version: 0.10.0.dev20210502

**Additional context**
Perhaps this is related to https://github.com/wandb/client/issues/2140


Discussion:
Yep, this is indeed related to the other ticket.  I'm working on [a PR ](https://github.com/wandb/client/pull/2150) to add a way to override this behavior that should land in the next release.
If I understand correctly, that PR adds an option to use `register_backward_hook` instead of `register_full_backward_hook`?
I actually believe using `register_full_backward_hook` is the correct thing to do, seeing as `register_backward_hook` is a deprecated function.

Note that on wandb version 0.10.27 the reproduction code produces the following warning:
```
UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.
```

I don't think `register_full_backward_hook` is the cause of the original `RuntimeError`, but rather the logic of the hook registered by `wandb.watch`.

Edit: Actually, notice that the reproduction code doesn't even perform a backward pass, and the `RuntimeError` occurs during the forward pass. I'm a bit baffled as to why the error (or warning in version 0.10.27) even mentions the backward hook. Is it being executed for some reason?
I encountered this also, downgrading to torchvision 0.8.2 solved the issue for me, in case thats of any help

[Pytorch forums thread](https://discuss.pytorch.org/t/runtimeerror-module-backward-hook-for-grad-input-is-called-before-the-grad-output-one-this-happens-because-the-gradient-in-your-nn-m-odule-flows-to-the-modules-input-without-passing-through-the-modules-output/119763/8)



Resolving patch:
diff --git a/wandb/wandb_torch.py b/wandb/wandb_torch.py
--- a/wandb/wandb_torch.py
+++ b/wandb/wandb_torch.py
@@ -311,6 +311,7 @@ def __init__(self):
         # started failing.  To maximize compatability we don't use full backward hooks
         # when we detect fastai v1 has been imported :(
         self._should_use_full_hooks = True
+        self._graph_hooks = {}
         if "fastai" in sys.modules:
             fastai = util.get_module("fastai")
             if fastai.__version__.startswith("1."):
@@ -322,13 +323,13 @@ def hook_torch(cls, model, criterion=None, graph_idx=0):
         graph.hook_torch_modules(model, criterion, graph_idx=graph_idx)
         return graph
 
-    def create_forward_hook(self, name, modules):
+    def create_forward_hook(self, name, graph_idx):
         graph = self
 
         def after_forward_hook(module, input, output):
-            if id(module) in modules:
+            if id(module) not in self._graph_hooks.keys():
+                # shound not happen
                 return
-            modules.add(id(module))
             if not isinstance(output, tuple):
                 output = (output,)
             parameters = [
@@ -354,12 +355,19 @@ def after_forward_hook(module, input, output):
                 elif isinstance(output[0], list) and hasattr(output[0][0], "grad_fn"):
                     graph.criterion = output[0][0].grad_fn
 
+            # log graph and remove hook
+            hook = self._graph_hooks.pop(id(module), None)
+            if hook is not None:
+                hook.remove()
+
+            if not self._graph_hooks:
+                # we went through the entire graph
+                wandb.run.summary["graph_%i" % graph_idx] = self
+
         return after_forward_hook
 
     def hook_torch_modules(self, module, criterion=None, prefix=None, graph_idx=0):
         torch = util.get_module("torch", "Could not import torch")
-        hooks = []
-        modules = set()
         layers = 0
         graph = self
         if hasattr(module, "_wandb_watch_called") and module._wandb_watch_called:
@@ -399,56 +407,9 @@ def hook_torch_modules(self, module, criterion=None, prefix=None, graph_idx=0):
             if isinstance(sub_module, tuple(module_types)):
                 self.hook_torch_modules(sub_module, prefix=name)
             else:
-
-                def backward_hook(module, input, output):
-                    [hook.remove() for hook in hooks]
-                    graph.loaded = True
-                    if wandb.run:
-                        wandb.run.summary["graph_%i" % graph_idx] = graph
-                    else:
-                        wandb.termwarn(
-                            "wandb.watch was called without a call to wandb.init, call wandb.init before wandb.watch",
-                            repeat=False,
-                        )
-                    # TODO: Keeping this here as a starting point for adding graph data
-                    if not graph.loaded:
-
-                        def traverse(node, functions=[]):
-                            if hasattr(node, "grad_fn"):
-                                node = node.grad_fn
-
-                            if hasattr(node, "variable"):
-                                node = graph.nodes_by_id.get(id(node.variable))
-                                if node:
-                                    node.functions = list(functions)
-                                    del functions[:]
-
-                            if hasattr(node, "next_functions"):
-                                functions.append(type(node).__name__)
-                                for f in node.next_functions:
-                                    if f[0]:
-                                        functions.append(type(f[0]).__name__)
-                                        traverse(f[0], functions)
-
-                            if hasattr(node, "saved_tensors"):
-                                for t in node.saved_tensors:
-                                    traverse(t)
-
-                        traverse(graph.criterion)
-
-                hooks.append(
-                    sub_module.register_forward_hook(
-                        self.create_forward_hook(name, modules)
-                    )
+                self._graph_hooks[id(sub_module)] = sub_module.register_forward_hook(
+                    self.create_forward_hook(name, graph_idx)
                 )
-                # Models with dicts as output must use register_full_backward_hook
-                if self._should_use_full_hooks and hasattr(
-                    sub_module, "register_full_backward_hook"
-                ):
-                    hook = sub_module.register_full_backward_hook(backward_hook)
-                else:
-                    hook = sub_module.register_backward_hook(backward_hook)
-                hooks.append(hook)
 
     @classmethod
     def from_torch_layers(cls, module_graph, variable):


Test patch: