Release trained model weights

[Nicholas Sofroniew]

Is there any chance @li-li-github @royerloic you could release trained weights for the PyTorch model that would allow me to use your autoencoder and the embedding space? Ideally it would be the same weights and umap parameters that allow me to remake figure 2 in your paper.

I’ve started trying to retrain the model with your data, but I’m not getting umaps that look as nice as yours in the paper. I can continue down that route, but if you have a nice trained model that you could add to the repo or hugging face that would be a big help!

thanks!!

[Hirofumi Kobayashi]

I don't have pretrained weights for the PyTorch implementation but I did check that PyTorch implementation works fairly similarly to our old Tensorflow implementation.
If you are using the simple_example.py, please try the following values.

model_args = {
...
    'fc_output_idx': 'all',
    'vq_args': {'num_embeddings': 2048, 'embedding_dim': 64},
    'fc_args': {'num_layers': 2},
...
}

train_args = {
    'lr': 1e-4,
    'max_epoch': 100,
    'reducelr_patience': 8,
    'reducelr_increment': 0.1,
    'earlystop_patience': 16,
}

Also, make sure to use all the data and train longer (at least > 30 epochs). The reconstruction loss 1 should be around 10^-3.

The model weights used in the paper can be found in the TensorFlow branch. You could try to convert the weights if that's easier than training a new one with the PyTorch implementation. (I'm sorry about this but I don't have a strong enough machine to run cytoself with the original setting right now so I can't validate the model or the implementation.)

For umap parameters, I basically used the default values but you can play with different n_neighbors and min_dist. Sometimes a bigger n_neighbors and a smaller min_dist could make the umap look slightly better.