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
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?