Package: autoencodersplz
A variety of autoencoder structured models for generative modeling and/or representation learning in pytorch. Models are mostly designed for usability/extensability/research rather than production implementations. But, go ahead and train some models and reconstruct some things!Table of contentsInstallationModelsLinearAELinearResidualAEConvResidualAEVQVAEFSQVAEMAEMAEMixIJEPATrainingBasicLightningExamplesBasic usageFuture additionsReferencesInstallationpipinstallautoencodersplzModelsLinearAEA fully-connected autoencoder with a linear/multi-layer perceptron encoder and decoderReducing the Dimensionality of Data with Neural Networksimporttorchfromautoencodersplz.modelsimportLinearAEmodel=LinearAE(img_size=224,in_chans=3,hidden_layers=[64,64],dropout_rate=0,latent_dim=16,beta=0.1,# beta > 0 = variationalmax_temperature=1000,# kld temperature annealingdevice=None)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)LinearResidualAEA fully-connected autoencoder with a linear/multi-layer perceptron residual network encoder and decoderSkip Connections Eliminate Singularitiesimporttorchfromautoencodersplz.modelsimportLinearResidualAEmodel=LinearResidualAE(img_size=224,in_chans=3,hidden_dim=[64,64],blocks=[2,2],dropout_rate=0.1,with_batch_norm=False,latent_dim=16,beta=0.1,# beta > 0 = variationalmax_temperature=1000,# kld temperature annealingdevice=None,)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)ConvResidualAEA convolutional autoencoder with a ResNet encoder and symmetric decoderDeep Residual Learning for Image Recognitionimporttorchfromautoencodersplz.modelsimportConvResidualAEmodel=ConvResidualAE(img_size=224,in_chans=3,channels=[64,128,256,512],blocks=[2,2,2,2],latent_dim=16,beta=0,# beta > 0 = variationalmax_temperature=1000,# kld temperature annealingupsample_mode='nearest',# interpolation methoddevice=None,)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)VQVAEA vector-quantized variational autoencoder with a ResNet encoder and symmetric decoderNeural Discrete Representation Learningimporttorchfromautoencodersplz.modelsimportVQVAEmodel=VQVAE(img_size=224,in_chans=3,channels=[64,128,256,512],blocks=[2,2,2,2],codebook_size=256,codebook_dim=8,use_cosine_sim=True,kmeans_init=True,commitment_weight=0.5,upsample_mode='nearest',vq_kwargs={},)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)FSQVAEA finite-scalar quantized variational autoencoder with a ResNet encoder and symmetric decoderFinite Scalar Quantization: VQ-VAE Made Simpleimporttorchfromautoencodersplz.modelsimportFSQVAEmodel=FSQVAE(img_size=224,in_chans=3,channels=[64,128,256,512],blocks=[2,2,2,2],levels=[8,6,5],upsample_mode='nearest')img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)MAEA masked autoencoder with a vision transformer encoder and decoderMasked Autoencoders Are Scalable Vision Learnersimporttorchimporttorch.nnasnnfromautoencodersplz.modelsimportMAEmodel=MAE(img_size=224,patch_size=16,in_chans=3,mask_ratio=0.5,embed_dim=768,depth=12,num_heads=12,mlp_ratio=4,pre_norm=False,decoder_embed_dim=768,decoder_depth=12,decoder_num_heads=12,norm_layer=torch.nn.LayerNorm,patch_norm_layer=torch.nn.LayerNorm,post_norm_layer=torch.nn.LayerNorm,)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)MAEMixA masked autoencoder with a MLP-mixer encoder and decoderMLP-Mixer: An all-MLP Architecture for Visionimporttorchfromautoencodersplz.modelsimportMAEMixmodel=MAEMix(img_size=224,patch_size=16,in_chans=3,mask_ratio=0.5,embed_dim=768,depth=12,mlp_ratio=4,decoder_embed_dim=768,decoder_depth=12,)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)IJEPAImage-based joint-embedding predictive architecture (Thanks toYiranfor porting this implementation)Self-Supervised Learning from Images with a Joint-Embedding Predictive Architectureimporttorchfromautoencodersplz.modelsimportIJEPAmodel=IJEPA(img_size=224,patch_size=16,in_chans=3,embed_dim=768,depth=12,num_heads=12,mlp_ratio=4,embed_dim_predictor=384,predictor_depth=12,num_targets=4,target_aspect_ratio=0.75,target_scale=0.2,context_aspect_ratio=1.,context_scale=0.9)img=torch.rand(1,3,224,224)loss,reconstructed_img=model(img)TrainingBasicTheTrainerclass enables basic training using a single CPU or GPU for any model in theautoencodersplzlibrary. TheTrainerclass will also automatically save the autoencoder model, backbone/encoder, losses, and a visualization of the training process (.gif) if you provide a path to theoutput_dirargument.fromautoencodersplz.trainersimportTrainertrainer=Trainer(autoencoder,train=train_dataloader,valid=valid_dataloader,epochs=128,learning_rate=5e-4,betas=(0.9,0.95),weight_decay=0.05,patience=10,scheduler='plateau',save_backbone=True,show_plots=False,output_dir='training_run/',device=None,)trainer.fit()By default,Traineruses anAdamWoptimizer and either aCosineDecay('cosine') orReduceLROnPlateau('plateau') scheduler. If you want to use different optimizers or schedulers, just re-assign a new optimizer or scheduler to the.optimizeror.schedulerattributes (withtrainer.model.parameters()) prior to callingtrainer.fit().LightningTo make it easier to scale to multi-gpu/distributed training, allautoencodersplzmodels are configured for use withpytorch lightning. Each model is setup with a default optimizer and scheduler and can be directly called by the pytorch lightning trainer. See an example below.importlightning.pytorchasplfromautoencodersplz.modelsimportFSQVAEmodel=FSQVAE(img_size=28,in_chans=1,channels=[8,16],blocks=[1,1],levels=[8],upsample_mode='nearest'learning_rate=1e-3,factor=0.1,patience=30,min_lr=1e-6)trainer=pl.Trainer(gpus=4,max_epochs=256)trainer.fit(model,train_dataloader,valid_dataloader)ExamplesBasic usageHere's a basic example of training a fully connected autoencoder on MNIST. The data is downloaded and loaded and then the autoencoder is fit. The training info is logged to the output directory (training/) and a GIF of the training routine is generated for visual inspection.fromtorch.utils.dataimportDataLoaderfromtorchvision.datasetsimportMNISTfromtorchvision.transformsimportToTensorfromautoencodersplz.modelsimportLinearAEfromautoencodersplz.trainersimportTrainertrain_loader=DataLoader(MNIST(root='data/',train=True,download=True,transform=ToTensor()),batch_size=32,shuffle=True,)test_loader=DataLoader(MNIST(root='data/',train=False,download=True,transform=ToTensor()),batch_size=32,shuffle=False,)model=LinearAE(img_size=28,in_chans=1,hidden_layers=[256,128],dropout_rate=0,latent_dim=32,beta=0,)trainer=Trainer(model,train_loader,test_loader,epochs=32,learning_rate=1e-3,output_dir='training/')trainer.fit()Future additionsEvolved Part Masking for Self-Supervised LearningRethinking Reconstruction Autoencoder-Based Out-of-Distribution DetectionCatch Missing Details: Image Reconstruction with Frequency Augmented
Variational AutoencoderDual Contradistinctive Generative AutoencoderScale-MAE: A Scale-Aware Masked Autoencoder for Multiscale Geospatial
Representation LearningGuided Variational Autoencoder for Disentanglement LearningMAGE: MAsked Generative Encoder to Unify Representation Learning and Image SynthesisReferences@article{hinton2006reducing,title={Reducing the dimensionality of data with neural networks},author={Geoffrey Hinton and Ruslan Salakhutdinov},url={10.1126/science.1127647},year={2006},}@article{orhan2018skip,title={Skip Connections Eliminate Singularities},author={Emin Orhan and Xaq Pitkow},url={https://arxiv.org/abs/1701.09175},year={2018},}@article{he2015deep,title={Deep Residual Learning for Image Recognition},author={Kaiming He and Xiangyu Zhang and Shaoqing Ren and Jian Sun},url={https://arxiv.org/abs/1512.03385},year={2016},}@misc{oord2018neural,title={Neural Discrete Representation Learning},author={Aaron van den Oord and Oriol Vinyals and Koray Kavukcuoglu},url={https://arxiv.org/abs/1711.00937},year={2017},}@misc{mentzer2023finite,title={Finite Scalar Quantization: VQ-VAE Made Simple},author={Fabian Mentzer and David Minnen and Eirikur Agustsson and Michael Tschannen},url={https://arxiv.org/abs/2309.15505},year={2023},}@misc{he2021masked,title={Masked Autoencoders Are Scalable Vision Learners},author={Kaiming He and Xinlei Chen and Saining Xie and Yanghao Li and Piotr Dollár and Ross Girshick},url={https://arxiv.org/abs/2111.06377},year={2021},}@misc{tolstikhin2021mlpmixer,title={MLP-Mixer: An all-MLP Architecture for Vision},author={Ilya Tolstikhin and Neil Houlsby and Alexander Kolesnikov and Lucas Beyer and Xiaohua Zhai and Thomas Unterthiner and Jessica Yung and Andreas Steiner and Daniel Keysers and Jakob Uszkoreit and Mario Lucic and Alexey Dosovitskiy},url={https://arxiv.org/abs/2105.01601},year={2021},}@misc{assran2023selfsupervised,title={Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture},author={Mahmoud Assran and Quentin Duval and Ishan Misra and Piotr Bojanowski and Pascal Vincent and Michael Rabbat and Yann LeCun and Nicolas Ballas},url={https://arxiv.org/abs/2301.08243},year={2023},}