Skip to content

Is it possible to partially run a model? #2998

@jamesheald

Description

@jamesheald

I have a module/class of this kind:

class autoencoder(nn.Module):
    
    hidden_dim: int
    z_dim: int
    output_dim: int

    def setup(self):
        
        self.encoder = encoder(self.hidden_dim, self.z_dim)
        self.decoder = decoder(self.hidden_dim, self.output_dim)

    def __call__(self, x):

        z = self.encoder(x)
        y = self.decoder(z)
        
        return y

Here, encoder and decoder are classes defined elsewhere. On occasions, I want to be able to call the individual submodules (e.g. self.encoder or self.decoder) of the module autoencoder, but as I understand they are inaccessible and can only be accessed by calling __ call __ (e.g. via model.apply), which calls both of them, consecutively (which I do also want to do sometimes). I ideally want to keep these submodules bundled together in this module so that when I initialize the model (via model.init), I have the full model and all of it's parameters in one place. (In my actual problem I have 6 submodules so I have simplified things here.)

This use case seems to be described here under 'Future work'. Is there no current solution to what I am trying to do?

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions