Reset model and optimizer in C++

Zhewen_Hu: <p>I am using CUDA graph to train multiple models. I believe the graph captures certain memory spaces, so if I need to train another network using the same graph, I have to use the same memory space of the model and optimizer. My question is how can I create a new model/optimizer in the same memory space, or is there a way to reset the model and the optimizer without changing the memory space?</p>
<p>Here is how I created the model and optimizer:<br>
std::shared_ptr&lt; UNet&gt; unet_model = std::make_shared&lt; UNet&gt;(in_chn, out_chn, fea);</p>
<p>std::shared_ptr&lt; torch::optim::Adam&gt; optimizer = std::make_shared&lt; torch::optim::Adam&gt;(unet_model-&gt;parameters(), torch::optim::AdamOptions(1e-2));</p>

ptrblck: <aside class="quote no-group" data-username="Zhewen_Hu" data-post="1" data-topic="213237">
<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/zhewen_hu/48/73939_2.png" class="avatar"> Zhewen_Hu:</div>
<blockquote>
<p>I believe the graph captures certain memory spaces, so if I need to train another network using the same graph, I have to use the same memory space of the model and optimizer</p>
</blockquote>
</aside>
<p>I don’t understand what this exactly means. The graph capture will use its own internal memory pool. You can capture any models you want as long as you have enough memory. If previous models are not needed anymore you can delete them.</p>

Zhewen_Hu: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,  I have a question about using a recorded training graph. If I record the training process of model_0 to build a CUDA graph, can I reuse the same graph to train model_1?</p>
<p>From my understanding, the graph seems to record the memory addresses associated with model_0, so during replay, it always references those specific addresses. To train model_1, would I need to copy model_1 to the same GPU memory address used by model_0 for the graph to work? Is this understanding correct?</p>

ptrblck: <aside class="quote no-group" data-username="Zhewen_Hu" data-post="3" data-topic="213237">
<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/zhewen_hu/48/73939_2.png" class="avatar"> Zhewen_Hu:</div>
<blockquote>
<p>If I record the training process of model_0 to build a CUDA graph, can I reuse the same graph to train model_1?</p>
</blockquote>
</aside>
<p>No, as your explanation about the captured memory addresses is correct. Replaying the graph will use the same operations and the same memory locations pointing to <code>model_0</code>.</p>
<aside class="quote no-group" data-username="Zhewen_Hu" data-post="3" data-topic="213237">
<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/zhewen_hu/48/73939_2.png" class="avatar"> Zhewen_Hu:</div>
<blockquote>
<p>To train model_1, would I need to copy model_1 to the same GPU memory address used by model_0 for the graph to work?</p>
</blockquote>
</aside>
<p>A safe option would be to delete <code>model_0</code> including its capture and to recapture <code>model_1</code>. However, you could also try to experiment with inplace copies (via <code>tensor.copy_</code>) to copy all parameters from <code>model_1</code> into <code>model_0</code> (I haven’t tried it and don’t know if I’m missing any obvious limitation).</p>

Zhewen_Hu: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="213237">
<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>However, you could also try to experiment with inplace copies (via <code>tensor.copy_</code>) to copy all parameters from <code>model_1</code> into <code>model_0</code> (I haven’t tried it and don’t know if I’m missing any obvious limitation).</p>
</blockquote>
</aside>
<p>Thanks, I forgot to mention that <code>model_0</code> and <code>model_1</code> share the same architecture, the only difference is their weights. Currently, I am training 4096 identical models, each with different weights. Since each model is relatively small (just a few hundred parameters), I believe reusing the graph could significantly improve performance, as capturing graphs for 4096 models might cause a lot of overhead.</p>
<p>My initial thought was to reuse the GPU address by creating a new model in the captured memory location, then after training copy it elsewhere, and then load the next model to repeat the process. However, if this approach isn’t feasible, I will explore the idea of copying tensors as you suggested.</p>

ptrblck: <p>If the architecture is equal, try to simply load the <code>state_dict</code> of another model into the captured one and let me know if you are seeing any issues.</p>