Using nn.utils.prune causes torch.Tensor.__deepcopy__ to fail

atila: <p>When pruning a module using utilities in the <code>torch.nn.utils.prune</code> (following the official <a href="https://pytorch.org/tutorials/intermediate/pruning_tutorial.html" rel="noopener nofollow ugc">PyTorch Pruning tutorial</a>), the “pruned” module becomes non-deep-copiable</p>
<p>Code to reproduce:</p>
<pre><code class="lang-auto">import copy
import torch
import torch.nn.utils.prune as prune

foo = torch.nn.Conv2d(2, 4, 1)
foo2 = copy.deepcopy(foo) # copy is successful before pruning

foo = prune.l1_unstructured(foo, name='weight', amount=0.2)
foo3 = copy.deepcopy(foo) # copy fails after pruning
</code></pre>
<p>The last line throws the following: <code>RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment</code>. because of <a href="https://github.com/pytorch/pytorch/blob/master/torch/tensor.py#L47" rel="noopener nofollow ugc">this check</a>.</p>
<p>Quick explanation of what the pruning call does:<br>
1-) It reregisters the pruned parameter with the “_orig” suffix <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/prune.py#L161" rel="noopener nofollow ugc">here</a><br>
2-) It registers the pruning mask as a buffer using the “_mask” suffix <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/prune.py#L177" rel="noopener nofollow ugc">here</a><br>
3-) It reuses the original weight name with <code>setattr(module, self._tensor_name, self.apply_mask(module))</code> <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/prune.py#L32" rel="noopener nofollow ugc">here</a></p>
<p>Only the last line seems to put a tensor that is not a leaf on the module:</p>
<pre><code class="lang-auto">print(foo.weight_orig.is_leaf) # True
print(foo.weight_mask.is_leaf) # True
print(foo.weight.is_leaf) # False
</code></pre>
<p>My 10 second solution was to monkey-patch <code>torch.Tensor.__deepcopy__</code> so that I can force a deepcopy by setting <code>self.is_leaf=True</code> and retry the deepcopy if it failed. My assumption was that forcing <code>is_leaf=True</code> has to do with properly copying gradients and minding the gradient tape. Since I intend to use the result of the deepcopy only for validation, I didn’t care about anything but forward pass correctness. (Edit: setting <code>is_leaf</code> is not allowed: <code>AttributeError: attribute 'is_leaf' of 'torch._C._TensorBase' objects is not writable</code>)</p>
<p>I am looking to understand the original motivation for the <a href="https://github.com/pytorch/pytorch/blob/master/torch/tensor.py#L47" rel="noopener nofollow ugc">is_leaf check in torch.Tensor.<strong>deepcopy</strong></a> so that I can see whether my assumption was correct.</p>
<p>Finally, a safer but tedious alternative is to remove the pruning reparameterization just before deepcopy so that all the modifications done by the pruner on the module are removed and the pruning mask is fused into the original weights. This works but it means that pruning is not resumable after removing the parameterization.</p>
<p>It would be awesome if someone (especially <a class="mention" href="/u/ptrblck">@ptrblck</a>) could propose a workaround or point out a mistake of mine to resolve this matter. Thanks in advance folks!</p>
<p>P.S. Not sure whether this is related yet but an <a href="https://discuss.pytorch.org/t/how-does-one-reproduce-the-deepcopy-error-of-tensors-made-by-code/69852/3">earlier topic</a> by <a class="mention" href="/u/pinocchio">@pinocchio</a> was talking about issues with <code>deepcopy</code> and using <code>setattr</code> to set tensors as attributes of the module. <code>setattr</code> is used several times in <code>torch.nn.utils.prune</code> (e.g. <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/prune.py#L32" rel="noopener nofollow ugc">here</a>). It might just have to do with  something else entirely.</p>

atila: <p>My current solution is to prepare the pruned module for a successful deepcopy by deleting the non-leaf tensor as an attribute of the module (since it is neither a parameter nor a buffer of the module)</p>
<pre><code class="lang-auto">import copy
import torch
import torch.nn.utils.prune as prune

foo = torch.nn.Conv2d(2, 4, 1)
foo2 = copy.deepcopy(foo) # copy is successful before pruning

pruner = prune.L1Unstructured.apply(foo, name='weight', amount=0.2)
foo3 = copy.deepcopy(foo) # copy fails after pruning

print(pruner._tensor_name) # 'weight'
delattr(foo, pruner._tensor_name)
foo4 = copy.deepcopy(foo) # copy is successful

print(hasattr(foo4, pruner._tensor_name)) # False
# Run forward() once to recover the deleted attribute 
_ = foo4(torch.empty(1, 2, 4, 4))
print(hasattr(foo4, pruner._tensor_name)) # True
</code></pre>
<p>Note that the deleted attribute is reset during the first <code>forward</code> call on the module. This is because it is being reregistered at every call in the <code>forward_pre_hook</code> that the <code>BasePruningMethod</code> class implements <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/prune.py#L32" rel="noopener nofollow ugc">here</a></p>
<p>cc: <a class="mention" href="/u/michela">@Michela</a>, any thoughts would be much appreciated for a longer term solution <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>