I have a pytorch model (custom model inherited from nn.Module).

I'm developing some architecture, for which makes sense for my task to have a list defined in the model as:

class Model(nn.Module):
   def __init__(self, hparams):
       self.list = []

In every training step, I want to store in this list some information. Particularly, a number that I will calculate over the training step.

To be more clear:

for each training step I receive a batch of size $n$, and in the training step I update the self.list in Model to store $n$ values. For instance, after a first iteration with a batch size of 100, the self.list list will contain 100 values. Every iteration, I will add 100 more to the list...

I mean, this makes sense for my task, but I'm wondering if this is "legal". Also it seems a bit weird, because the parameters of my model would increase with the training size? Because in the end, if I use 10,000 training samples, this list will hold 10,000 values. But if I use more, the list will have more. Would this be an increase in parameters? When I save this model, will the number of parameters be affected by this? That's what I'm not sure about.


1 Answer 1


Generally, your PyTorch module is to contain the architecture that you are training, not the additional variables that you want to store. So generally, I would not implement my list like this.

What is usually done is to have one nn.Module for the architecture and one generic class for utility functions such as the model training/fitting/validation etc. Then the module is passed to the generic 'Trainer' class for training the model and your list of values would presumably be stored in the 'Trainer' object.

So something like this:

class MyNetwork(nn.Module):
    # Here define the architecture
    def __init__(self, *other_arguments): ...
    def forward(batch): ...

class Trainer():
    def __init__(self, network_instance, optimizer, *other_arguments):
        self.network_instance = network_instance
        self.optimizer = optimizer

        self.list_of_values = []

    def train(n_epochs):
        for i in range(n_epochs):
            # Standard train function layout
            batch = self.get_data()
            predictions = self.network_instance(batch)
            loss = calculate_loss(predictions)

            # Add value to list
  • $\begingroup$ I understand. It indeed makes sense to store the list in the Trainer. Now, if I store a separate model (let's say an sklearn model) in a variable of the nn.Module, and I use this accumulated list to train that model, (calling it in the trainer as sklearn_model.fit(list_of_values), will those learned weights be succesfully stored inside the model variable in the nn.Module after training? $\endgroup$ Apr 13 at 13:41
  • $\begingroup$ I am not familiar with the integration of sklearn models (and their parameters) into PyTorch Modules. I'd advise keeping them separate if possible. $\endgroup$ Apr 13 at 14:13
  • $\begingroup$ I actually think they are not made to work together in terms of saving sklearn model weights for instance, at least from my experiments. I guess the right choice would be to implement the sklearn model in Pytorch and then load this model as an instance of the nn.Module class. Many thanks. $\endgroup$ Apr 13 at 14:42
  • $\begingroup$ I think thats the right approach indeed. If the answer sufficiently answers your question make sire to upvote it and select it as the correct answer $\endgroup$ Apr 13 at 19:35

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .