Initialization of register_buffer for loading model

thomas: <p>Hello,</p>
<p>I am in a situation were I initialize and train a model using some datasets. I then save the model state dict. I would like to be able to load the model without having to load the data again. However because of the dimensions mismatch when I call <code>load_state_dict</code> I don’t seem to be able to do that.</p>
<p>example :</p>
<pre><code class="lang-auto">class Model(nn.Module):
    def __init__(
            self,
            layer_list,
            x_mean: Optional[np.array] = None,
            x_std: Optional[np.array] = None,
            **_
    ):
        super(Model, self).__init__()

        self.register_buffer('x_mean', None if x_mean is None else torch.FloatTensor(x_mean))
        self.register_buffer('x_std', None if x_std is None else torch.FloatTensor(x_std))
        self.sequential = get_sequential(layer_list)
</code></pre>
<p>After training the <code>x_mean</code>, <code>x_std</code>, … are filled with tensor of certain shapes. However when I load the model, I ‘don’t know’ the shape of those, and I would like to avoid to search in the data what are the shape of x_mean and x_std.</p>
<p>In this case I would like to load that model, just using the <code>layer_list</code> and the function <code>load_state_dict</code></p>
<p>Thanks!</p>

ptrblck: <p>The <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer"><code>register_buffer</code></a> operation includes a <code>persistent</code> argument, which defines if the buffer should be added to the <code>state_dict</code> (it’s set to <code>True</code> by default). If you don’t want to store <code>x_mean</code> and <code>x_std</code> in the <code>state_dict</code> you could thus set it to <code>False</code> for these buffers.</p>

thomas: <p>Actually, I would like to save it in the state dict. However when I initialize the model (before loading the state dict), I do not know the shape of those buffers. Therefore I got an error “shape mismatch” if I initialize it with a random tensor or just None if I initialize them with None. I would like to be able to load the model without explicitly have to say the shape of the buffers, so when it loads it can assign whatever tensors to those buffers.</p>
<p>Minimal working example:</p>
<pre><code class="lang-auto">import torch 
from typing import Optional

class Model(torch.nn.Module):
    def __init__(self,x: Optional[list] = None):
        super(Model, self).__init__()
        self.register_buffer('x', None if x is None else torch.FloatTensor(x))

PATH = 'state_dict.pt'
x = [1,  2]

model = Model(x)
torch.save(model.state_dict(), PATH)

model_loaded = Model()
model_loaded.load_state_dict(torch.load(PATH))
print(model_loaded.x)
</code></pre>

ptrblck: <p>In your code snippet you are loading the <code>state_dict</code> into <code>model</code>, while I assume you want to load it into <code>model_loaded</code>, which will raise a <code>RuntimeError</code>, since no <code>x</code> buffer was registered:</p>
<pre><code class="lang-python">import torch 
from typing import Optional

class Model(torch.nn.Module):
    def __init__(self,x: Optional[list] = None):
        super(Model, self).__init__()
        self.register_buffer('x', None if x is None else torch.FloatTensor(x))

PATH = 'state_dict.pt'
x = [1,  2]

model = Model(x)
sd = model.state_dict()

model_loaded = Model()
model_loaded.load_state_dict(sd)
&gt; RuntimeError: Error(s) in loading state_dict for Model:
	Unexpected key(s) in state_dict: "x". 
</code></pre>
<p>The proper way would be to make sure to create this buffer, since you already know the shape from the loaded <code>state_dict</code>:</p>
<pre><code class="lang-python">model_loaded = Model(sd['x'])
model_loaded.load_state_dict(sd)
&gt; &lt;All keys matched successfully&gt;
</code></pre>

thomas: <p>Oh, ok<br>
This is exactly what I needed !<br>
Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>(Yes I did a typo in the minimal working example)</p>