How to remove in-place copy_ operator from the traced graph?

apivovarov: <p>I traced one of the <a href="https://github.com/facebookresearch/detectron2" rel="noopener nofollow ugc">detectron2</a> models and noticed that the resulting graph has in-place <code>copy_</code> operators.</p>
<p>Is it possible to somehow remove/replace this <code>copy_</code> operator from the graph?<br>
Maybe <code>jit</code> has an option to produce a graph without in-place operators (such as copy_)?</p>

apivovarov: <p>Probably it is better to rewrite the following 4 lines in the model itself.<br>
Detectron2 - <a href="https://github.com/facebookresearch/detectron2/blob/3c6857a9ba7ccb14eac5265955526366662f9da9/detectron2/modeling/box_regression.py#L111" rel="noopener nofollow ugc">box2box_transform.apply_deltas:111</a></p>

ptrblck: <p>Changing these lines of code might remove the <code>copy_</code> ops and you could try it by locally creating this model with the changes.<br>
However, just out of interest, why would you like to remove these operations from the graph?</p>

tom: <p>I think the background is moving the model to TVM see <a href="https://github.com/apache/tvm/pull/7513" rel="noopener nofollow ugc">this discussion on TVM and <code>copy_</code></a>.</p>
<p><a class="mention" href="/u/apivovarov">@apivovarov</a> You probably increase your chance of getting good answers if you give more context.</p>
<p>Personally, I think the two main ways of doing this is:</p>
<ul>
<li>As you suggest, the immediate way is to change  the model to use <code>torch.cat</code> + view instead. This is probably wildly preferable because any translation of index assignment look funny.</li>
<li>Implement view tracking in TVM. I commented a bit on this in the TVM <code>copy_</code> PR before yours. Takes some effort but after that it should be reasonably possible to express index assignments. There will likely be corner cases that you miss as there are operations PyTorch (e.g. <code>reshape</code>) that may or may not return views.</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

eellison: <p>I think probably the best way to go about this is to prove that the input is linearly typed and replace <code>copy_</code> with a functional op that creates a copy of the tensor and then writes to it. This is basically how other functional tensor languages handle writes - like <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_update.html" rel="noopener nofollow ugc">jax</a> and <a href="https://www.cl.cam.ac.uk/~nk480/numlin.pdf" rel="noopener nofollow ugc">numlin</a>. You need to prove that the input is linearly typed (there are no other aliases to it) so that you don’t change the semantics of the program.</p>
<p>I maintain a <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/remove_mutation.cpp" rel="noopener nofollow ugc">pass</a> that does this logic for many inplace ops but does not yet handle <code>copy_</code>. I would accept a PR that adds a new operator similar to <code>jax.ops.index_update</code> and replaces <code>copy_</code> with it when that can be proved safe.</p>
<p>Looking at <a href="https://github.com/facebookresearch/detectron2/blob/3c6857a9ba7ccb14eac5265955526366662f9da9/detectron2/modeling/box_regression.py#L110" rel="noopener nofollow ugc">apply_deltas</a>, we initialize the tensor and then write to it so we would successively be able to remove <code>copy_</code> from this function.</p>

apivovarov: <p>Another place which probably should be rewritten in Detectron2 is <code>boxes.clip()</code> - <a href="https://github.com/facebookresearch/detectron2/blob/master/detectron2/structures/boxes.py#L205" class="inline-onebox" rel="noopener nofollow ugc">detectron2/boxes.py at master · facebookresearch/detectron2 · GitHub</a></p>