Variational Transformer Autoencoder with Manifolds Learning
This is a PyTorch implementation for our paper "Variational Transformer Autoencoder with Manifolds Learning". It is an open-source codebase for image interpolations and reconstructions.
We test this repo with Python 3.8, PyTorch 1.9.0, and CUDA 11.1. However, it should be runnable with recent PyTorch versions (Pytorch >= 1.1.0).
python setup.py developIn addition, we need to install KNN_CUDA on the wheel.
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whlThe file tree we used for storing the pre-trained weights is like
logs
├── pretrained.pth.tar(1) Celeba-pretrained weights for CNNs backbone
The Celeba-pretrained weights for CNN's backbone or the pretrained weights for the model.
Train by running a script in the terminal. train.py
In the train.py. In case you want to speed up testing, enlarge GPUS for more GPUs, or enlarge the --tuple-size for more tuples on one GPU. In case your GPU does not have enough memory, reduce --pos-num or --neg-num for fewer positives or negatives in one tuple.
Test by running a script in the terminal. Script test.py
If you find our work useful for your research, please consider citing our paper:
@article{shamsolmoali2023vtae,
title={VTAE: Variational Transformer Autoencoder with Manifolds Learning},
author={Shamsolmoali, Pourya and Zareapoor, Masoumeh and Zhou, Huiyu and Tao, Dacheng and Li, Xuelong},
journal={IEEE Transactions on Image Processing},
year={2023}
}