Saving TensorDictModule

putskan: <p>Hey!<br>
What is the recommended approach for saving a trained model to disk in TorchRL?</p>
<p>As opposed to standard Torch, saving <code>nn.Module</code> directly might not be the best way (for example, in cases where we chain using <code>TensorDictSequential</code>).</p>
<p>Much appreciated!</p>

Kwool: <p>Hey,</p>
<p>I am new here, but wanted to try and help out. I threw your question into Claude Sonnet 3.5 (My teacher). Sonnet recommended a few methods but <strong>.state_dict()</strong> seems like the winner.</p>
<ol>
<li>Saving the State Dict: The most flexible and recommended way is to save the state dictionary of your model. This works well even with complex TorchRL structures.</li>
</ol>
<p>To save:<br>
torch.save(model.state_dict(), “model_state_dict.pth”)</p>
<p>To load:<br>
model.load_state_dict(torch.load(“model_state_dict.pth”))</p>
<p>Note: .state_dict() should work well for TensorDictSequential.<br>
Ex: torch.save(tensor_dict_sequential.state_dict(), “tensordict_sequential_state.pth”)</p>
<ol start="2">
<li>Using TensorDict: If your model uses TensorDict structures, you might want to save the entire TensorDict:</li>
</ol>
<p>from tensordict import TensorDict</p>
<p>To save:<br>
Assuming your model state is in a TensorDict,<br>
state_dict = TensorDict({“model”: model.state_dict()})<br>
torch.save(state_dict, “model_tensordict.pth”)</p>
<p>To load:<br>
loaded_state = torch.load(“model_tensordict.pth”)<br>
model.load_state_dict(loaded_state[“model”])</p>
<ol start="3">
<li>Saving the Entire Model: While not always recommended, especially for complex structures, you can save the entire model:</li>
</ol>
<p>torch.save(model, “full_model.pth”)</p>
<p>Load with:<br>
model = torch.load(“full_model.pth”)</p>
<p>I look forward to seeing what others recommend, and hope this helps.</p>
<p>Regards,<br>
Kyle</p>

putskan: <p>Thanks for the answer mate,<br>
You’re right. Didn’t think it’s so well supported.</p>
<p>For the sake of completeness, below is a code that verifies it:</p>
<pre><code class="lang-auto">import copy

import torch
from tensordict.nn import TensorDictModule, TensorDictSequential
from torch import nn


def save_checkpoint_using_state_dict(module: TensorDictModule, dst_path: str) -&gt; None:
    torch.save(module.state_dict(), dst_path)


def load_checkpoint_using_state_dict(module: TensorDictModule, src_path: str) -&gt; TensorDictModule:
    module.load_state_dict(torch.load(src_path))
    return module


def save_checkpoint(module: TensorDictModule, dst_path: str) -&gt; None:
    torch.save(module, dst_path)


def load_checkpoint(src_path: str) -&gt; TensorDictModule:
    return torch.load(src_path)


if __name__ == '__main__':
    path = '/tmp/model.pth'
    m1 = nn.Linear(1, 1)
    m2 = nn.Linear(1, 1)
    module1 = TensorDictModule(m1, in_keys='inp', out_keys='out')
    module2 = TensorDictModule(m2, in_keys='out', out_keys='out2')
    module = TensorDictSequential(module1, module2)
    module_copy = copy.deepcopy(module)

    # using state dict:
    save_checkpoint_using_state_dict(module, path)
    with torch.no_grad():
        [p.zero_() for p in module.parameters()]

    module = load_checkpoint_using_state_dict(module, path)
    for p1, p2 in zip(module.parameters(), module_copy.parameters()):
        assert torch.equal(p1.data, p2.data)

    # no state dict:
    save_checkpoint(module, path)
    with torch.no_grad():
        [p.zero_() for p in module.parameters()]

    module = load_checkpoint(path)
    for p1, p2 in zip(module.parameters(), module_copy.parameters()):
        assert torch.equal(p1.data, p2.data)
</code></pre>
<p>Cheers!</p>