Accessing part of a model (eg. encoder of an Autoencoder) when defined using setup

[nipunbatra]

I have an Autoencoder module defined below.

class Encoder(nn.Module):
    bottleneck: int
        
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(5)(x)
        x = nn.selu(x)
        x = nn.Dense(features=self.bottleneck)(x)
        return x

class Decoder(nn.Module): 
    out: int
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(5)(x)
        x = nn.selu(x)
        x = nn.Dense(features=self.out)(x)
        return x

class AE(nn.Module):
    bottleneck: int
    out: int
    def setup(self):
        self.encoder = Encoder(bottleneck=self.bottleneck)
        self.decoder = Decoder(out=self.out)

    def __call__(self, x):
        
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat

ae = AE(1, 2)
params = ae.init(random.PRNGKey(0), X)

When I try to access the ae.encoder, I get the following error.

ae.encoder

gives the error

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [100], in <cell line: 1>()
----> 1 ae.encoder

File ~/miniconda3/lib/python3.8/site-packages/flax/linen/module.py:783, in Module.__getattr__(self, name)
    780 if self.scope is None:
    781   msg += (f' If "{name}" is defined in \'.setup()\', remember these fields '
    782     'are only accessible from inside \'init\' or \'apply\'.')
--> 783 raise AttributeError(msg)

AttributeError: "AE" object has no attribute "encoder". If "encoder" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'.

As a workaround, I am able to do the following
encoded_1d = Encoder(1).apply({"params":params["params"]["encoder"]}, X)

What could be a simpler/better way to access the encoder part of the AutoEncoder module and use a .apply on it with the encoder parameters?

[cgarciae]

Hey @nipunbatra, there are 2 ways to go about this.

Define an API

You can define methods to interact with specific parts of you model and use them via the method argument e.g:

class AE(nn.Module):
    bottleneck: int
    out: int
    def setup(self):
        self.encoder = Encoder(bottleneck=self.bottleneck)
        self.decoder = Decoder(out=self.out)

    def __call__(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat
        
    def encode(self, x):
        return self.encoder(x)
        
    def decode(self, z):
        return self.decoder(z)
        
ae = AE(1, 2)
variables = ae.init(random.PRNGKey(0), X)
params = variables['params']

# call encoder
z = ae.apply({'params': params}, X, method=ae.encode)

# call decoder
x_hat = ae.apply({'params': params}, X, method=ae.decode)

Extract submodules

You can use nn.apply with some helper functions to extract submodules defined inside .setup and their corresponding variables.

def extract_encoder(ae):
    encoder_module = ae.encoder.clone() # avoid leakage
    encoder_params =  ae.encoder.variables['params']
    return encoder_module, encoder_params
    
def extract_decoder(ae):
    decoder_module = ae.decoder.clone() # avoid leakage
    decoder_params =  ae.decoder.variables['params']
    return decoder_module, decoder_params
    
encoder, encoder_params = nn.apply(extract_encoder, ae)({'params': params})
decoder, decoder_params = nn.apply(extract_decoder, ae)({'params': params})

# call encoder
z = encoder.apply({'params': encoder_params}, X)

# call decoder
x_hat = decoder.apply({'params': decoder_params}, z)