How to save nn.Sequential as a model?

Fernando_Gallego: <p>Hi, I am trying to decompose ResNet into three different devices, for this, I would need to be able to save their nn.sequential as a separate model. Training them all together but being able to load their models separately on each device.</p>
<p>Do any of you know how to save nn.sequential  as a model? I 'm working as every morning</p>

iffiX: <p>if you won’t change your model-device mapping, you can just save your model directly using <code>t.save</code> and load it with <code>t.load</code>.<br>
If you really want to save <code>nn.sequential</code>, you can also save it direcly using <code>t.save</code> and load it with <code>t.load</code>.</p>

Fernando_Gallego: <p>But if, for instance, I have a neural network with a structure like this:</p>
<p>layers1 = nn.sequential(…)<br>
layers2 = nn.sequential(…)<br>
layers3 = nn.sequential(…)</p>
<p>And I would like to save the model so that device1 loads layers1, device2 layers2…</p>

iffiX: <p>for example, suppose you initialize a model like:</p>
<pre><code class="lang-auto">class YourModel(nn.Module):
    def __init__(self, dev_list=["cpu", "cuda:1", "cuda:0"]):
        self.fc1 = nn.sequential(nn.Linear(5, 5).to(dev_list[0]))  # on device cpu
        self.fc2 = nn.sequential(nn.Linear(5, 5).to(dev_list[1]))  # on device "cuda:1"
        self.fc3 = nn.sequential(nn.Linear(5, 5).to(dev_list[2]))  # on device "cuda:0"
        self.dev = dev_list

    def forward(x):
        x = self.fc1(x).to(self.dev[1])
        x = self.fc2(x).to(self.dev[1])
        return self.fc3(x).to("cpu")
</code></pre>
<p>then:</p>
<pre><code class="lang-auto">t.save(YourModel(), "model.pt")
model = t.load("model.pt")
</code></pre>
<p>Device mapping will be saved along with your model, don’t worry about it</p>

Fernando_Gallego: <p>So, if I wanted to map only fc1 on device 1, could I select it on the load?</p>

iffiX: <p>See <a>https://pytorch.org/docs/stable/torch.html?highlight=load#torch.load</a>, especially the <code>map_location</code> part</p>

Fernando_Gallego: <p>But in this solution, the devices must be connected and assigned. My goal is to deploy the neural network, in a distributed way, decomposing each sequential in different devices and sending the inference through the network.</p>

iffiX: <p>Could you please clarify your design a liitle bit more? your description of “distributed” and “decomposing” is pretty vague.</p>
<p>Is it a multi-process application? How do you map your gpus to your processes? How do you map your model to your devices? I am really sorry that I cannot help you more if you don’t <strong>give a clear definition of the architecture</strong> you would like to achieve. It would be better if you can draw a graph or show your code.</p>
<p>If you just want to split <code>layer1</code>, <code>layer2</code> and <code>layer3</code> to different devices, you can simply save them individually with <code>torch.save</code> and <code>torch.load</code>, torch will take the charge to pickle whatever passed to it, including parameters and your custom attributes such as the ones set by <code>self.some_attr</code> in <code>__init__</code>.</p>

Fernando_Gallego: <p>It’s an architecture, fog, edge and cloud. Based on the example above:</p>
<p>class YourModel(nn.Module):<br>
def <strong>init</strong>(self):<br>
self.fc1 = nn.sequential()<br>
self.fc2 = nn.sequential()<br>
self.fc3 = nn.sequential()<br>
I intend to save fc1, fc2 and fc3 separately. In this way I could make a first prediction in the device that has the model fc1 and sending the inference to the second device, make the next prediction with greater accuracy. The third one would work in the same way.</p>

iffiX: <p>I see, one more question, will you move the model around, such as on a different machine with different gpu number, or are you loading the whole model on the same devices?</p>
<p>If you don’t, and you really want to save them seperately to different files, maybe for better inspection or archive perpose, then:</p>
<pre><code class="lang-auto">def save(your_model):
    torch.save(your_model.fc1, "fc1.pt")
    torch.save(your_model.fc2, "fc2.pt")
    torch.save(your_model.fc3, "fc3.pt")
</code></pre>
<p>If you do, then you will have to decide which device each part of your model would locate on, eg: suppose on your training machine you have 3 gpus, and on your inference machine you have 1 gpu.<br>
def save(your_model):</p>
<pre><code class="lang-auto">def save(your_model):
    torch.save(your_model.fc1, "fc1.pt")
    torch.save(your_model.fc2, "fc2.pt")
    torch.save(your_model.fc3, "fc3.pt")

def map(your_model):
    your_model.fc1 = torch.load("fc1.pt", map_location=torch.device('cuda:0'))
    your_model.fc2 = torch.load("fc2.pt", map_location=torch.device('cuda:0'))
    your_model.fc3 = torch.load("fc3.pt", map_location=torch.device('cuda:0'))
</code></pre>
<p>by the way,</p>
<aside class="quote no-group" data-username="Fernando_Gallego" data-post="8" data-topic="89117">
<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/fernando_gallego/48/26683_2.png" class="avatar"> Fernando_Gallego:</div>
<blockquote>
<p>the devices must be connected and assigned</p>
</blockquote>
</aside>
<p>Maybe you have some wrong idea, there is not such a “connected device” concept in pytorch, you can perform a complex forward() operation or a simple add() operation on some input <code>x</code> locating on device <code>cuda:[number]</code> or <code>cpu</code> simply because the operands (tensors) locates on the same device, if torch needs to fetch it somewhere else, it will complain and throw an error.</p>
<h2><a name="about-saving-the-model-1" class="anchor" href="#about-saving-the-model-1"></a>About saving the model</h2>
<p>There are many ways to save your model, typically you will want to save the <code>OrderedDict</code> returned by <code>model.state_dict()</code>, the keys are your parameter names such as “linear.weight” or “linear.bias”, and values are <code>nn.Parameter</code>, its <code>.data</code> attribute is just a Tensor. You may load a state dict into your model like:</p>
<pre><code class="lang-auto">def prep_load_state_dict(model: nn.Module,
                         state_dict: Any):
    """
    Automatically load a **loaded state dictionary**

    Note:
        This function handles tensor device remapping.
    """
    for name, param in model.named_parameters():
        state_dict[name].to(param.device)
    model.load_state_dict(state_dict)
</code></pre>
<h2><a name="about-torchsave-and-torchload-2" class="anchor" href="#about-torchsave-and-torchload-2"></a>About torch.save and torch.load</h2>
<p>If you know the <code>pickle</code> concept in python, then you will get what <code>torch.save</code> does. <code>pickle</code> serialize a object into binary string:</p>
<pre><code class="lang-auto">buffer = io.BytesIO()
t.save(t.zeros([5]), buffer)
print(buffer.getvalue())
</code></pre>
<p>will yield:</p>
<pre><code class="lang-auto">b'\x80\x02\x8a\nl\xfc\x9cF\xf9 j\xa8P\x19.\x80\x02M\xe9\x03.\x80\x02}q\x00(X\n\x00\x00\x00type_sizesq\x01}q\x02(X\x03\x00\x00\x00intq\x03K\x04X\x04\x00\x00\x00longq\x04K\x04X\x05\x00\x00\x00shortq\x......
</code></pre>
<p>you can serialize whatever you like into this, cuda tensor will essentially be saved as “raw data” + “device descriptor cuda:0”.</p>

Fernando_Gallego: <aside class="quote no-group" data-username="iffiX" data-post="11" data-topic="89117">
<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/iffix/48/24443_2.png" class="avatar"> iffiX:</div>
<blockquote>
<pre><code class="lang-auto">def save(your_model):
    torch.save(your_model.fc1, "fc1.pt")
    torch.save(your_model.fc2, "fc2.pt")
    torch.save(your_model.fc3, "fc3.pt")
</code></pre>
</blockquote>
</aside>
<p>Thanks a lot, I think the solution were:</p>
<p>def save(your_model):<br>
torch.save(your_model.fc1, “fc1.pt”)<br>
torch.save(your_model.fc2, “fc2.pt”)<br>
torch.save(your_model.fc3, “fc3.pt”)</p>
<p>I’ll try then.</p>

iffiX: <p>Great! post your issues if you have any</p>

Brando_Miranda: <p>I’m also trying to do something similar but in my scenario I construct a whole model using only <code>nn.Sequential</code> and then I just want to save it. I don’t have a class defined for it so something like <a href="https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch" rel="nofollow noopener">https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch</a> won’t work for me.</p>
<hr>
<p>My current attempt uses pickle but I keep getting warning for using pickle:</p>
<pre><code class="lang-auto">FutureWarning: pickle support for Storage will be removed in 1.5. Use `torch.save` instead
  warnings.warn("pickle support for Storage will be removed in 1.5. Use `torch.save` instead", FutureWarning)
</code></pre>
<hr>
<p>I think they just want us to use <code>torch.save</code> and <code>torch.load</code>. I stopped getting warning when I did that.</p>
<p>My (full) code:</p>
<pre><code class="lang-auto">
# creating data and running through a nn and saving it

import torch
import torch.nn as nn

from pathlib import Path
from collections import OrderedDict

import numpy as np

import pickle

path = Path('~/data/tmp/').expanduser()
path.mkdir(parents=True, exist_ok=True)

num_samples = 3
Din, Dout = 1, 1
lb, ub = -1, 1

x = torch.torch.distributions.Uniform(low=lb, high=ub).sample((num_samples, Din))

f = nn.Sequential(OrderedDict([
    ('f1', nn.Linear(Din,Dout)),
    ('out', nn.SELU())
]))
y = f(x)

# save data torch to numpy
x_np, y_np = x.detach().cpu().numpy(), y.detach().cpu().numpy()
np.savez(path / 'db', x=x_np, y=y_np)

print(x_np)
# save model
with open('db_saving_seq', 'wb') as file:
    pickle.dump({'f': f}, file)

# load model
with open('db_saving_seq', 'rb') as file:
    db = pickle.load(file)
    f2 = db['f']

# test that it outputs the right thing
y2 = f2(x)

y_eq_y2 = y == y2
print(y_eq_y2)

db2 = {'f': f, 'x': x, 'y': y}
torch.save(db2, path / 'db_f_x_y')

print('Done')

db3 = torch.load(path / 'db_f_x_y')
f3 = db3['f']
x3 = db3['x']
y3 = db3['y']
yy3 = f3(x3)

y_eq_y3 = y == y3
print(y_eq_y3)

y_eq_yy3 = y == yy3
print(y_eq_yy3)
</code></pre>
<p>did you try that? Is there a reason why that’s not enough for you?</p>
<hr>
<p>cross-posted: <a href="https://stackoverflow.com/questions/62923052/how-does-one-save-torch-nn-sequential-models-in-pytorch-properly" rel="nofollow noopener">https://stackoverflow.com/questions/62923052/how-does-one-save-torch-nn-sequential-models-in-pytorch-properly</a></p>