Package: vision-aided-loss
Vision-aided GANvideo|website|paperCan the collectiveknowledgefrom a large bank of pretrained vision models be leveraged to improve GAN training? If so, with so many models to choose from, which one(s) should be selected, and in what manner are they most effective?We find that pretrained computer vision models can significantly improve performance when used in an ensemble of discriminators.  We propose an effective selection mechanism, by probing the linear separability between real and fake samples in pretrained model embeddings, choosing the most accurate model, and progressively adding it to the discriminator ensemble. Our method can improve GAN training in both limited data and large-scale settings.Ensembling Off-the-shelf Models for GAN TrainingNupur Kumari,Richard Zhang,Eli Shechtman,Jun-Yan ZhuIn CVPR 2022Pretrained ModelsOur final trained models can be downloaded at thislink. For more details on usage please see README in the folderstylegan2andbiggan.Vision-aided StyleGAN2 trainingPlease seestylegan2README for training StyleGAN2 models with our method. This code will reproduce all StyleGAN2 based results from our paper.Vision-aided Discriminator in a custom GAN modelpip install vision-aided-lossinstalling from github:pip install git+https://github.com/nupurkmr9/vision-aided-gan.gitorgit clone https://github.com/nupurkmr9/vision-aided-gan.git
cd vision-aided-gan
pip install .For details on off-the-shelf models please seeMODELS.mdimportvision_aided_lossdevice='cuda'discr=vision_aided_loss.Discriminator(cv_type='clip',loss_type='multilevel_sigmoid_s',device=device).to(device)discr.cv_ensemble.requires_grad_(False)# Freeze feature extractor# Sample imagesreal=sample_real_image()fake=G.forward(z)# Update discriminator discrlossD=discr(real,for_real=True)+discr(fake,for_real=False)lossD.backward()# Update generator GlossG=discr(fake,for_G=True)lossG.backward()# We recommend adding vision-aided adversarial loss after training GAN with standard loss till few warmup_iter.Arg details:cv_type: name of the off-the-shelf model from[clip, dino, swin, vgg, det_coco, seg_ade, face_seg, face_normals]. Multiple models can be used with '+' separated model names.output_type: output feature type from off-the-shelf models. should be one of[conv, conv_multi_level]. Supportsconv_multi_levelonly for clip and dino. For multiple models output_type should be '+' separated output_type for each model.diffaug: if True performs DiffAugment on vision-aided discriminator with poilcycolor,translation,cutout. Recommended to keep this as True.num_classes: for conditional training use num_classes>0. Projection discriminator is used similar toBigGAN.loss_type: should be one of[sigmoid, multilevel_sigmoid, sigmoid_s, multilevl_sigmoid_s, hinge, multilevel_hinge]. Appeding_senableslabel smoothing. If loss_type is None output is a list of logits corresponding to each vision-aided discriminator.device: device for off-the-shelf model weights.Vision-aided StyleGAN3 trainingPlease seestylegan3README for training StyleGAN3 models with our method.Vision-aided BigGAN trainingPlease seebigganREADME for training BigGAN models with our method.To add you own pretrained Modelcreate the class file to extract pretrained features asvision_module/<custom_model>.py. Add the class path in theclass_name_dictinvision_module.cvmodel.CVBackboneclass. Update the architecture of trainable classifier head over pretrained features invision_module.cv_discriminator. Reinstall library viapip install .References@InProceedings{kumari2021ensembling,
  title={Ensembling Off-the-shelf Models for GAN Training},
  author={Kumari, Nupur and Zhang, Richard and Shechtman, Eli and Zhu, Jun-Yan},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  month     = {June},
  year      = {2022}
}AcknowledgmentsWe thank Muyang Li, Sheng-Yu Wang, Chonghyuk (Andrew) Song for proofreading the draft. We are also grateful to Alexei A. Efros, Sheng-Yu Wang, Taesung Park, and William Peebles for helpful comments and discussion. Our codebase is built onstylegan2-ada-pytorchandDiffAugment.