# Use pre-trained autoencoder for classification or regression

**URL:** <https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147>\
**Category:** vision\
**Created:** [January 27, 2021, 3:20pm UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147 "2021-01-27T15:20:48Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![anotherone\_one](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/anotherone_one/32/18256_2.png) [@anotherone\_one](https://discuss.pytorch.org/u/anotherone_one)\
**Post date:** [January 27, 2021, 3:20pm UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147/1 "2021-01-27T15:20:48Z")

</div>

Hello!!

I trained an autoencoder and now I want to use that model with the trained weights for classification purposes. So, I suppose I have to freeze the weights and layer of the encoder and then add classification layers, but I am a bit confused on how to to this.

This is the AutoEncoder I trained

```auto
class AE(nn.Module):
    def __init__ (self, **kwargs):
        super(). __init__ ()
        self.encoder_hidden_layer = nn.Linear(
            in_features=kwargs["input_shape"], out_features=128
        )
        self.encoder_output_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_hidden_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_output_layer = nn.Linear(
            in_features=128, out_features=kwargs["input_shape"]
        )

    def forward(self, features):
        activation = self.encoder_hidden_layer(features)
        activation = torch.relu(activation)
        code = self.encoder_output_layer(activation)
        code = torch.relu(code)
        activation = self.decoder_hidden_layer(code)
        activation = torch.relu(activation)
        activation = self.decoder_output_layer(activation)
        reconstructed = torch.relu(activation)
        return reconstructed

```

Can someone help me?

Thank you

---

<div class="post-metadata">

**Author:** ![Dwight\_Foster](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dwight_foster/32/18182_2.png) [@Dwight\_Foster](https://discuss.pytorch.org/u/Dwight_Foster)\
**Post date:** [January 27, 2021, 6:33pm UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147/2 "2021-01-27T18:33:35Z")

</div>

You could create a separate model then pass the autoencoder into it and load the weights in that. Then you could create a separate sequential layer with the classification layers and then in the forward pass you could pass the input through those two. Then too freeze the autoencoder weights in the optimizer you could just set the parameters as only the classification layers.

---

<div class="post-metadata">

**Author:** ![anotherone\_one](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/anotherone_one/32/18256_2.png) [@anotherone\_one](https://discuss.pytorch.org/u/anotherone_one)\
**Post date:** [January 27, 2021, 6:35pm UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147/3 "2021-01-27T18:35:27Z")

</div>

Can you show me in code how to do that?

---

<div class="post-metadata">

**Author:** ![Dwight\_Foster](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dwight_foster/32/18182_2.png) [@Dwight\_Foster](https://discuss.pytorch.org/u/Dwight_Foster)\
**Post date:** [January 27, 2021, 7:14pm UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147/4 "2021-01-27T19:14:59Z")

</div>

Yes here is the model code you could use

```auto
class classification(nn.Module):
    def __init__ (self, **kwargs):
        super(). __init__ ()
        self.AE = AE(kwargs)
        self.AE.load_state_dict(FILENAME)
        self.classifier = nn.Sequential(nn.Linear(INPUT,OUPUT))
        ##You can use as many linear layers and other activations as you want
    def forward(self, input):
        x = self.AE(input)
        output = self.classifier(output)
        ##You can use a sigmoid or softmax here if neecd
        return output

```

and then the optimizer

```auto
optimizer = torch.optim.Adam(model.classifier.parameters(), lr=LR)

```

---

<div class="post-metadata">

**Author:** ![anotherone\_one](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/anotherone_one/32/18256_2.png) [@anotherone\_one](https://discuss.pytorch.org/u/anotherone_one)\
**Post date:** [January 28, 2021, 9:20am UTC](https://discuss.pytorch.org/t/use-pre-trained-autoencoder-for-classification-or-regression/110147/5 "2021-01-28T09:20:43Z")

</div>

Thank you so much!!!