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

1

Dear all,
I have some model Class that are defined as below.

class Net(nn.Module):    
    def __init__(self) -> None:
        super(Net, self).__init__()
     <...>

    def forward(self, x: torch.Tensor) -> torch.Tensor:
   <...>
        return x

I have multiple model classes, so use some config dict file to define what model is to be actually instanced.

if cfg.model = 'NET'
    class Model(Net):
        pass
elif cfg.model = 'NET1'
    class Model(Net1):
        pass

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…

model = models.efficientnet_b0(pretrained=pretrained)
2

One of the setups that I had come across is here:

You can look at the __init__.py file to see if this design would help you.

3

@InnovArul Thanks for the suggestion, I am not sure if this is the right approach for me…

I think my current problem is, I am trying to inherit from two kinds of base classes.
My previous model are all from nn.Module:

class Net(nn.Module):  
<...>
# and I can use it like below:
UsedModel = Net().to(device)

Now I have a model builder as below:

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

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.
what I did is like below:

class TLNet():
    def __init__(self, pretrained=True, input_, output):
        self = build_model(pretrained, input_, output)

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…

# for this code: 
UsedModel = TLNet().to(device), 
# I get the error as below:
AttributeError: 'UsedModel' object has no attribute 'to'

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.

Can anyone give advice?

4

Maybe I did not understand the whole context.
My question is, what stops TLNet() to inherit from nn.Module?

What is the motivation to assign self here?

5

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.

Then I just make another change like the below:

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

Then I was able to pass the Model.to(cuda), but when running the model training, I get the cuda error report like:

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 >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

This error should come from my change as if I use my older model there is no such error report.
Also, if do not use such a Class instance, but directly generate a model to use, send to cuda and train, it also works…

6

A straightforward reason is that you are indeed using more batch-size or more GPU memory somehow.
I do not see technically why a class instance uses more memory than a direct model.

7

Thanks, I find the problem, that I was using an incorrect transformer – now my program fully works.
Thank you a lot for the suggestion!