Best practices for saving/loading dynamic pytorch models

ua01: <p>Hi all,</p>
<p>I am currently working on a project for which I require to simulate the behaviour of some device using neural networks. I would like to have some flexibility in the architecture that is used for simulating the device, so I have created a function that allows me to do that inside the model:</p>
<pre><code class="lang-auto">class DeviceModel(nn.Module):
    def __init__(self, custom_dictionary):
        super(DeviceModel, self).__init__()
        self.build_my_structure(custom_dictionary) 

    def build_my_structure(self, custom_dictionary):
        hidden_sizes = custom_dictionary["hidden_sizes"]
        input_layer = nn.Linear(custom_dictionary["D_in"], hidden_sizes[0])
        activation = self._get_activation(custom_dictionary["activation"])   
        output_layer = nn.Linear(hidden_sizes[-1], custom_dictionary["D_out"])        
        modules = [input_layer, activation]

        hidden_layers = zip(hidden_sizes[:-1], hidden_sizes[1:])
        for h_1, h_2 in hidden_layers:
            hidden_layer = nn.Linear(h_1, h_2)
            modules.append(hidden_layer)
            modules.append(activation)

        modules.append(output_layer)
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        return self.model(x)
</code></pre>
<p>What would be the best practice here to load the model? Should I save it directly using torch.save, or should I save the state_dictionary. If the second way is the best, I have to create the model from scratch using the original configs dictionary and then load the state_dict, how should can I store both of them together?</p>
<p>Perhaps, I am just using the library wrong, and there are other ways of creating a model dynamically? Or there is simply no support for this?</p>
<p>Thank you very much in advance. Any help would be much appreciated.</p>

ptrblck: <aside class="quote no-group" data-username="ua01" data-post="1" data-topic="115603">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/u/e274bd/48.png" class="avatar"> ua01:</div>
<blockquote>
<p>What would be the best practice here to load the model? Should I save it directly using torch.save, or should I save the state_dictionary.</p>
</blockquote>
</aside>
<p>I would recommend to save and load the <code>state_dict</code> and to recreate the model using a stored config, as the former approach might break in various ways e.g. if you change the folder structure etc.</p>
<aside class="quote no-group" data-username="ua01" data-post="1" data-topic="115603">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/u/e274bd/48.png" class="avatar"> ua01:</div>
<blockquote>
<p>If the second way is the best, I have to create the model from scratch using the original configs dictionary and then load the state_dict, how should can I store both of them together?</p>
</blockquote>
</aside>
<p>Yes, you would have to recreate the model. You can store the <code>state_dict</code> as well as other objects e.g. in a <code>dict</code> and use <code>torch.save()</code> on it as shown in the <a href="https://github.com/pytorch/examples/blob/af111380839d35924e9e36437aeb9757b5a68f96/imagenet/main.py#L256-L262">ImageNet example</a>.</p>