Encoder implementation for image inversion task of stylegan3 generator (Alias Free GAN).
The neural network architecture and hyper-parameter settings of the base configuration is almost the same as that of pixel2style2pixel, and various settings of improved encoder architecture will be added in the future.
For fast training, pytorch DistibutedDataParallel is used.
Please see this repo for further research (Stylegan3-edit).
- GeForce RTX 3090 x 8
- NVIDIA driver version: 460.91.03
$ sh build_img.sh
$ sh build_container.sh [container-name]
$ docker start [container-name]
$ docker attach [container-name]
$ pip install -v -e .
python train.py \
--outdir exp/[exp_name] \
--encoder [encoder_type] \
--data data/[dataset_name] \
--gpus [num_gpus] \
--batch [total_batch_size] \
--generator [generator_pkl]
python test.py \
--testdir exp/[train_exp]/[train_exp_subdir] \
--data data/[dataset_name] \
--gpus [num_gpus] \
--batch [total_batch_size]
Train options
{
"model_architecture": "base",
"dataset_dir": "data/ffhq",
"num_gpus": 8,
"batch_size": 32,
"batch_gpu": 4,
"generator_pkl": "pretrained/stylegan3-t-ffhq-1024x1024.pkl",
"val_dataset_dir": null,
"training_steps": 100001,
"val_steps": 10000,
"print_steps": 50,
"tensorboard_steps": 50,
"image_snapshot_steps": 100,
"network_snapshot_steps": 5000,
"learning_rate": 0.001,
"l2_lambda": 1.0,
"lpips_lambda": 0.8,
"id_lambda": 0.1,
"reg_lambda": 0.0,
"gan_lambda": 0.0,
"edit_lambda": 0.0,
"random_seed": 0,
"num_workers": 3,
"resume_pkl": null,
"run_dir": "exp/base/00000-base-ffhq-gpus8-batch32"
}
Trainset examples
Real image batch X
Encoded image batch G.synthesis(E(X))
Testset examples(celeba-hq)
Target image
Encoded image
Encoded image, transform x=0.2, y=0
Encoded image, transform x=0.2, y=0.1
Encoded image, transform x=-0.2, y=0.1
Encoded image, transform x=-0.2, y=-0.1