Load part of big model

username: mmilakov
cooked: <p>Hi,</p>
<p>I have a model saved with <code>torch.save</code> into the file. The file is quite big (say, 100 GB), <code>torch.load</code> is crashing on a nodes with moderate CPU RAM. The good part is that I don’t need all the tensors in CPU memory at once.</p>
<p>Is there a way to customize <code>torch.load</code> so that it doesn’t produce model with all tensors deserialized? Basically, the majority of tensors should be dropped right after each of them being deserialized or not deserialized at all.</p>
<p><a href="https://pytorch.org/docs/stable/torch.html#torch.load" rel="nofollow noopener">The docs suggest</a> <code>map_location</code> callable argument could possibly do the trick, I cannot find examples for that though.</p>
<p>Any help is appreciated!</p>
<p>Thanks,<br>
Max.</p>

username: ptrblck
cooked: <p>I assume you are using the proper approach of storing the <code>state_dict</code> instead of the complete model.<br>
If that’s the case, you could remove the keys you don’t want to store to disk.<br>
After loading this smaller <code>state_dict</code> from disk again, load it to your model with <code>strict=False</code>:</p>
<pre><code class="lang-python">model = nn.Sequential(
    nn.Linear(1, 1),
    nn.Linear(1, 2))

sd = model.state_dict()

# Delete unnecessary keys
del sd['0.weight']
del sd['0.bias']

# Store
torch.save(sd, 'tmp.pt')

# Load
model = nn.Sequential(
    nn.Linear(1, 1),
    nn.Linear(1, 2))
sd = torch.load('tmp.pt')
model.load_state_dict(sd, strict=False)
</code></pre>

username: mmilakov
cooked: <p>Thank you for your answer! I have to add I don’t have control over saving the model. My input is a file with the saved model (you are right, this is state_dict).</p>

username: ptrblck
cooked: <p>Oh, that’s trickier.<br>
If I’m not mistaken, PyTorch uses <code>pickle</code> to store and load the data.<br>
It might be possible to load only part of the file via <code>pickle</code>, e.g. by using an <code>io.BytesIO</code> object and <code>seek</code> manually.<br>
However, this sounds like a cumbersome approach and you would also need to dive into the <a href="https://github.com/pytorch/pytorch/blob/master/torch/serialization.py">serialization code</a>.</p>
<p>Could you load the <code>state_dict</code> once on a machine with enough memory and then store parts of it to use the first approach?</p>

username: mmilakov
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="74892">
<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>Could you load the <code>state_dict</code> once on a machine with enough memory and then store parts of it to use the first approach?</p>
</blockquote>
</aside>
<p>This is the workaround I am currently implementing <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
Thanks for your help, I will probably use this workaround for some time, things might change (for example, I might get some control over model saving code).</p>