Modified implementation of AGE for few-shot image generation. Code is modified from AGE.

• removed the seed fixing during inference. Now one seed is used to generate one image. All seeds are preset.
• added guassian noise in the datapipe line
• added orthogonal loss for global dictionary A

• Linux
• NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported)
• Python 3

Follow the instructions to train a pSp model firsly. Or you can also directly download the pSp pre-trained models.

• organize the file structure as follows:

  └── data_root
      ├── train                      
      |   ├── cate-id_sample-id.jpg                # train-img
      |   └── ...                                  # ...
      └── valid                      
          ├── cate-id_sample-id.jpg                # valid-img
          └── ...                                  # ...
  

The format of the file should be [label_id]_[any_name].jpg

Path	Description
AGE	Repository root folder
├  configs	Folder containing configs defining model/data paths and data transforms
├  criteria	Folder containing various loss criterias for training
├  datasets	Folder with various dataset objects and augmentations
├  environment	Folder containing Anaconda environment used in our experiments
├ models	Folder containting all the models and training objects
│  ├  encoders	Folder containing our pSp encoder architecture implementation and ArcFace encoder implementation from TreB1eN
│  ├  stylegan2	StyleGAN2 model from rosinality
│  └  age.py	Implementation of AGE
├  options	Folder with training and test command-line options
├  tools	Folder with running scripts for training and inference
├  optimizer	Folder with Ranger implementation from lessw2020
└  utils	Folder with various utility functions

• Refer to configs/paths_config.py to define the necessary data paths and model paths for training and evaluation.
• Refer to configs/transforms_config.py for the transforms defined for each dataset.
• Finally, refer to configs/data_configs.py for the data paths for the train and valid sets as well as the transforms.
• To experiment with your own dataset, you can simply make the necessary adjustments in
  1. data_configs.py to define your data paths.
  2. transforms_configs.py to define your own data transforms.

To train AGE, the class embedding of each category in training set should be get first by using tools/get_class_embedding.py.

cd AGE; python tools/get_class_embedding.py \
--class_embedding_path=/path/to/save/classs/embeddings \
--psp_checkpoint_path=/path/to/pretrained/pSp/checkpoint \
--train_data_path=/path/to/training/data \
--test_batch_size=4 \
--test_workers=4

The main training script can be found in tools/train.py.
Intermediate training results are saved to opts.exp_dir. This includes checkpoints, train outputs, and test outputs.
Additionally, if you have tensorboard installed, you can visualize tensorboard logs in opts.exp_dir/logs.

#set GPUs to use.
export CUDA_VISIBLE_DEVICES=0,1,2,3

#begin training.
cd AGE; python -m torch.distributed.launch \
--nproc_per_node=4 \
tools/train.py \
--dataset_type=af_encode \
--exp_dir=/path/to/experiment/output \
--workers=8 \
--batch_size=8 \
--valid_batch_size=8 \
--valid_workers=8 \
--val_interval=2500 \
--save_interval=5000 \
--start_from_latent_avg \
--l2_lambda=1 \
--sparse_lambda=0.005 \
--orthogonal_lambda=0.0005 \
--A_length=100 \
--psp_checkpoint_path=/path/to/pretrained/pSp/checkpoint \
--class_embedding_path=/path/to/class/embeddings 

Having trained your model or using pre-trained models provided, you can use tools/inference.py to apply the model on a set of images.
For example,

cd AGE;python tools/inference.py \
--output_path=/path/to/output \
--checkpoint_path=/path/to/checkpoint \
--test_data_path=/path/to/test/input \
--train_data_path=/path/to/training/data \
--class_embedding_path=/path/to/classs/embeddings \
--n_distribution_path=/path/to/save/n/distribution \
--test_batch_size=4 \
--test_workers=4 \
--n_images=5 \
--alpha=1 \
--beta=0.005

specify different train_data_path and n_distribution_path for different emotion labels

cd AGE;python tools/inference.py \
--output_path=/path/to/output \
--checkpoint_path=/path/to/checkpoint \
--test_data_path=/path/to/test/input \
--train_data_path=/path/to/training/data/emotion02 \
--class_embedding_path=/path/to/classs/embeddings \
--n_distribution_path=/path/to/save/n/distribution/emotion02 \
--test_batch_size=4 \
--test_workers=4 \
--n_images=5 \
--alpha=1 \
--beta=0.005

@inproceedings{ding2022attribute,
  title={Attribute Group Editing for Reliable Few-shot Image Generation},
  author={Ding, Guanqi and Han, Xinzhe and Wang, Shuhui and Wu, Shuzhe and Jin, Xin and Tu, Dandan and Huang, Qingming},
  booktitle=CVPR,
  year={2022},
}