How to wrap a pre-trained model to a class?

Leonmac: <p>Dear all,<br>
I have some model Class that are defined as below.</p>
<pre><code class="lang-auto">class Net(nn.Module):    
    def __init__(self) -&gt; None:
        super(Net, self).__init__()
     &lt;...&gt;

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
   &lt;...&gt;
        return x
</code></pre>
<p>I have multiple model classes, so use some config dict file to define what model is to be actually instanced.</p>
<pre><code class="lang-auto">if cfg.model = 'NET'
    class Model(Net):
        pass
elif cfg.model = 'NET1'
    class Model(Net1):
        pass
</code></pre>
<p>Now if I want to use pre-trained models like below-- but at the same time I would like to warp that into a Class similar to the above “Model”, so I don’t need to change the older code. then how could I do that? Apparently, I can not inherit the model from nn.Module any more…</p>
<pre><code class="lang-auto">model = models.efficientnet_b0(pretrained=pretrained)
</code></pre>

InnovArul: <p>One of the setups that I had come across is here:</p><aside class="onebox githubfolder" data-onebox-src="https://github.com/KaiyangZhou/deep-person-reid/tree/master/torchreid/models">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/KaiyangZhou/deep-person-reid/tree/master/torchreid/models" target="_blank" rel="noopener">github.com</a>
  </header>

  <article class="onebox-body">
    <h3><a href="https://github.com/KaiyangZhou/deep-person-reid/tree/master/torchreid/models" target="_blank" rel="noopener">deep-person-reid/torchreid/models at master · KaiyangZhou/deep-person-reid</a></h3>

  <p><a href="https://github.com/KaiyangZhou/deep-person-reid/tree/master/torchreid/models" target="_blank" rel="noopener">master/torchreid/models</a></p>

  <p><span class="label1">Torchreid: Deep learning person re-identification in PyTorch.</span></p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>You can look at the <code>__init__.py</code> file to see if this design would help you.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/models/__init__.py">
  <header class="source">

      <a href="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/models/__init__.py" target="_blank" rel="noopener">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/models/__init__.py" target="_blank" rel="noopener">KaiyangZhou/deep-person-reid/blob/master/torchreid/models/__init__.py</a></h4>


      <pre><code class="lang-py">from __future__ import absolute_import
import torch

from .pcb import *
from .mlfn import *
from .hacnn import *
from .osnet import *
from .senet import *
from .mudeep import *
from .nasnet import *
from .resnet import *
from .densenet import *
from .xception import *
from .osnet_ain import *
from .resnetmid import *
from .shufflenet import *
from .squeezenet import *
from .inceptionv4 import *
from .mobilenetv2 import *
from .resnet_ibn_a import *
</code></pre>



  This file has been truncated. <a href="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/models/__init__.py" target="_blank" rel="noopener">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Leonmac: <p><a class="mention" href="/u/innovarul">@InnovArul</a> Thanks for the suggestion, I am not sure if this is the right approach for me…</p>
<p>I think my current problem is, I am trying to inherit from two kinds of base classes.<br>
My previous model are all from nn.Module:</p>
<pre><code class="lang-auto">class Net(nn.Module):  
&lt;...&gt;
# and I can use it like below:
UsedModel = Net().to(device)
</code></pre>
<p>Now I have a model builder as below:</p>
<pre><code class="lang-auto">def build_model(model_name, pretrain, input_, output):
    MODEL = getattr (models, model_name)
    model = MODEL(pretrained=pretrain)
    
    model.classifier[1] = nn.Linear(in_features=input_, out_features=output)
    return model
</code></pre>
<p>Ultimately, I want to wrap such a model builder into a class, so I can still inherit from one base class, based on the model defined in the config.<br>
what I did is like below:</p>
<pre><code class="lang-auto">class TLNet():
    def __init__(self, pretrained=True, input_, output):
        self = build_model(pretrained, input_, output)
</code></pre>
<p>but it does not work as the instance I get from the TLNet is not an object that can be used to pass to the device like cuda…</p>
<pre><code class="lang-auto"># for this code: 
UsedModel = TLNet().to(device), 
# I get the error as below:
AttributeError: 'UsedModel' object has no attribute 'to'
</code></pre>
<p>The problem is that the UsedModel is not inherited from a base class like nn.Module and it is not recognized as a valid model to pass to the device.</p>
<p>Can anyone give advice?</p>

InnovArul: <aside class="quote no-group" data-username="Leonmac" data-post="3" data-topic="152271">
<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/leonmac/48/30837_2.png" class="avatar"> Leonmac:</div>
<blockquote>
<p>The problem is that the UsedModel is not inherited from a base class like nn.Module and it is not recognized as a valid model to pass to the device.</p>
</blockquote>
</aside>
<p>Maybe I did not understand the whole context.<br>
My question is, what stops <code>TLNet()</code> to inherit from <code>nn.Module</code>?</p>
<aside class="quote no-group" data-username="Leonmac" data-post="3" data-topic="152271">
<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/leonmac/48/30837_2.png" class="avatar"> Leonmac:</div>
<blockquote>
<p><code>self = build_model(pretrained, input_, output)</code></p>
</blockquote>
</aside>
<p>What is the motivation to assign <code>self</code> here?</p>

Leonmac: <p>Actually, I don’t know if that is the correct thing to do–I just wanted to get an instance of the actual model that I need, so make a self=… that was ugly.</p>
<p>Then I just make another change like the below:</p>
<pre><code class="lang-auto">class TLNet(nn.Module):
    def __init__(self, net_name = "efficientnet_b0", pretrain=True, input_=1024, output=4):
        super(TLNet, self).__init__()

        MODEL = getattr (models, net_name)
        self.net = MODEL(pretrained=pretrain)

        self.net.classifier[1] = nn.Linear(in_features=input_, out_features=output)

    def forward(self, x: torch.Tensor):
        x = self.net(x)
        return x
</code></pre>
<p>Then I was able to pass the Model.to(cuda), but when running the model training, I get the cuda error report like:</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory. Tried to allocate 18.00 MiB (GPU 0; 3.95 GiB total capacity; 2.02 GiB already allocated; 44.44 MiB free; 2.05 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
</code></pre>
<p>This error should come from my change as if I use my older model there is no such error report.<br>
Also, if do not use such a Class instance, but directly generate a model to use, send to cuda and train, it also works…</p>

InnovArul: <aside class="quote no-group" data-username="Leonmac" data-post="5" data-topic="152271">
<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/leonmac/48/30837_2.png" class="avatar"> Leonmac:</div>
<blockquote>
<p><code>RuntimeError: CUDA out of memory.</code></p>
</blockquote>
</aside>
<p>A straightforward reason is that you are indeed using more batch-size or more GPU memory somehow.<br>
I do not see technically why a class instance uses more memory than a direct model.</p>

Leonmac: <p>Thanks, I find the problem, that I was using an incorrect transformer – now my program fully works.<br>
Thank you a lot for the suggestion!</p>