How to make a dataloader from a frozen network?

Anton1: <p>I’d like to use a frozen network as a type of preprocessing for a later network which is being trained. Instead of freezing layers, I’d like to combine my original train_dataloader and the frozen network into a new wrapper-dataloader. How should I best approach that?</p>
<p>I’ve written</p>
<pre><code class="lang-auto">class NetLoader:
    def __init__(self, loader, net):
        self.loader = loader
        self.net = net

    def __iter__(self):
        with torch.no_grad():
            for X, Y in self.loader:
                yield self.net(X), Y
</code></pre>
<p>and I use this as an iterable “dataloader” later when I train my proper network. I see multiple issues:</p>
<ul>
<li>it’s not a proper DataLoader class</li>
<li>I’m not sure if/where to use <code>torch.no_grad()</code>, <code>net.requires_grad=False</code> and/or <code>self.net(X).detach()</code>
</li>
</ul>
<p>What is the proper way to do that?</p>
<p>The main goal is to build a network layer-wise (instead of picking out parts from a full specification).</p>
<p>Any other suggestions?</p>

JuanFMontesinos: <p>Why don’t you just create a nn.Module that wraps both? You can paralellize that, you don’t have to worry about devices as a simple cuda call will affect all the submodules properly.<br>
This seems a headache to me.</p>
<p>Anyway, torch nograd is enough for the purpose.</p>

Anton1: <p>You mean something like</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, fixed_net, train_net):
        super().__init__()
        self.fixed_net = fixed_net
        self.train_net = train_net
        
    def forward(self, X):
        with torch.no_grad():
            X = self.fixed_net(X)
            
        return self.train_net(X)
    
    def parameters(self):
        return self.train_net.parameters()
</code></pre>
<p>?<br>
Am I missing something here?</p>
<p>So <code>no_grad</code> is enough and no future operation will somehow leak into my fixed net and perform unnecessary computation?</p>

JuanFMontesinos: <p>Indeed,<br>
that context manager disables autograd engine. See: <a href="https://pytorch.org/docs/stable/generated/torch.no_grad.html" class="inline-onebox" rel="noopener nofollow ugc">no_grad — PyTorch 1.12 documentation</a></p>
<p>You may also want to overwrite load_state_dict and state_dict if you are not interested in saving/loading both at the same time. This way, when you allocate net or run in parallel the network, everything will work smoothly.</p>

Anton1: <p>Thanks. I see. I try that.</p>
<p>For <code>no_grad</code>, I understand that it does what I want during the forward pass. I was concerned that unneeded computation could happened during later <code>optim.zero_grad()</code> or <code>out.backward()</code> or <code>optim.step()</code>.</p>

JuanFMontesinos: <p>Werll for optim.step it just computes the backward for those parameters you pass to the optimizer. in your case, as you are rewritting the method <code>parameters</code>, you will never pass those.<br>
Besides, there will not be graph to backprop, it won’t be computer or stored. Lastly, you would need a leaf node before the first network (and there will be no).</p>