Can I implement a sklearn model inside a Pytorch nn.Module? [closed]

Question

I am making a custom Pytorch model that at some point, clusters a latent space that was created by another, previous routine of the model (Autoencoder).

In a bit more detail, my model is a regular Autoencoder, but in every training step, I want to perform a clustering of the representations in the latent space Z, and save the clustering algorithm inside the model to use after training, ONLY for inference (to make certain kind of predictions).

I thought about defining an sklearn model as an instance of the nn.Module as

class Model(Module):
    def __init__(self, hparams):
        super().__init__()
        self.clustering_model = sklearn.clustering_model()

This way, I could call the clustering model to fit the latent space Z, after calculating the reconstruction loss as a regular autoencoder.

After training, I would then have in the variable model.clustering_model a trained clustering model that I could use for inference.

However, this doesn't seem to work out, probably because of incompatibilities between sklearn and Pytorch. Basically, it seems that training the clustering_model and saving it inside the pytorch model didn't actually save any trained weights.

I'm right now unsure whether it didn't save weights because it actually wasn't trained (by an error on my side) or if it is that PyTorch just can't save a sklearn trained model as a variable in a nn.Module.

Did anyone try before to implement sklearn inside a nn.Module?

Answer 1

Yes, you can define a totally custom Model, maybe with a clustering method you can call after forward only during inference. Clustering parameters (e.g centroids in Kmeans) can be stored inside a buffer (see here and here)

Clustering is not typically optimized with Gradient Descent, so clustering parameters should not require gradient, and you need to pay attention to not include the clustering step into your computational graph.

I would implement this pipeline:

1. Autoencoder training: Train your AutoEncoder as usual, and save the model's checkpoint.
2. Feature extraction: Extract all latent codes using your trained autoencoder, then save them onto disk into a numpy array or a tensor.
3. Clustering: Load all latent codes and apply a clustering algorithm (e.g Kmeans) and save all relevant parameters (centroids) onto a buffer .
4. Inference: Load trained autoencoder and centroids. Both parameters and buffers are into model's state_dict, so they can be loaded with the load_state_dict method. Pass new data into the encoder, then cluster the latent representation onto the closest centroid

By using a modular approach, you can experiment different clustering methods without losing all the features (saving additional training), and then put all together once you're satisfied with the model.