Simple Beta-VAE using scikit-learn API, made mostly by prompting GPT-4. With a single argument this can instead be a regular autoencoder (no variational). The VeLO optimizer can be used (apparently only on cpu and not on cuda?)

I made this because I couldn't find an appropriate implementation in python / pytorch and needed one for another project: QuestEA

from bvae import ReducedBVAE
model = ReducedBVAE(
    input_dim,
    z_dim,  # lowest nmuber of dimension
    hidden_dim,  # number of neurons in the 2nd layer of the compression
    dataset_size,
    lr=1e-3,
    epochs=1000,
    beta=1.0,
    weight_decay=0.01,
    use_VeLO=False,
    use_scheduler=True,
    )
model.prepare_dataset(
    dataset=dataset,
    val_ratio=0.2,
    batch_size=500,
)
model.train_bvae(
    patience=100,
)
projection = model.transform(dataset)