\n$ python3 context_encoder.py\n```\n\n\n ![](\"assets/context_encoder.png\")
\n
\n\n Rows: Masked | Inpainted | Original | Masked | Inpainted | Original\n
\n\n### Coupled GAN\n_Coupled Generative Adversarial Networks_\n\n#### Authors\nMing-Yu Liu, Oncel Tuzel\n\n#### Abstract\nWe propose coupled generative adversarial network (CoGAN) for learning a joint distribution of multi-domain images. In contrast to the existing approaches, which require tuples of corresponding images in different domains in the training set, CoGAN can learn a joint distribution without any tuple of corresponding images. It can learn a joint distribution with just samples drawn from the marginal distributions. This is achieved by enforcing a weight-sharing constraint that limits the network capacity and favors a joint distribution solution over a product of marginal distributions one. We apply CoGAN to several joint distribution learning tasks, including learning a joint distribution of color and depth images, and learning a joint distribution of face images with different attributes. For each task it successfully learns the joint distribution without any tuple of corresponding images. We also demonstrate its applications to domain adaptation and image transformation.\n\n[[Paper]](https://arxiv.org/abs/1606.07536) [[Code]](implementations/cogan/cogan.py)\n\n#### Run Example\n```\n$ cd implementations/cogan/\n$ python3 cogan.py\n```\n\n\n ![](\"assets/cogan.gif\")
\n
\n\n Generated MNIST and MNIST-M images\n
\n\n### CycleGAN\n_Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks_\n\n#### Authors\nJun-Yan Zhu, Taesung Park, Phillip Isola, Alexei A. Efros\n\n#### Abstract\nImage-to-image translation is a class of vision and graphics problems where the goal is to learn the mapping between an input image and an output image using a training set of aligned image pairs. However, for many tasks, paired training data will not be available. We present an approach for learning to translate an image from a source domain X to a target domain Y in the absence of paired examples. Our goal is to learn a mapping G:X→Y such that the distribution of images from G(X) is indistinguishable from the distribution Y using an adversarial loss. Because this mapping is highly under-constrained, we couple it with an inverse mapping F:Y→X and introduce a cycle consistency loss to push F(G(X))≈X (and vice versa). Qualitative results are presented on several tasks where paired training data does not exist, including collection style transfer, object transfiguration, season transfer, photo enhancement, etc. Quantitative comparisons against several prior methods demonstrate the superiority of our approach.\n\n[[Paper]](https://arxiv.org/abs/1703.10593) [[Code]](implementations/cyclegan/cyclegan.py)\n\n\n ![](\"http://eriklindernoren.se/images/cyclegan.png\")
\n
\n\n#### Run Example\n```\n$ cd data/\n$ bash download_cyclegan_dataset.sh monet2photo\n$ cd ../implementations/cyclegan/\n$ python3 cyclegan.py --dataset_name monet2photo\n```\n\n\n ![](\"assets/cyclegan.png\")
\n
\n\n Monet to photo translations.\n
\n\n### Deep Convolutional GAN\n_Deep Convolutional Generative Adversarial Network_\n\n#### Authors\nAlec Radford, Luke Metz, Soumith Chintala\n\n#### Abstract\nIn recent years, supervised learning with convolutional networks (CNNs) has seen huge adoption in computer vision applications. Comparatively, unsupervised learning with CNNs has received less attention. In this work we hope to help bridge the gap between the success of CNNs for supervised learning and unsupervised learning. We introduce a class of CNNs called deep convolutional generative adversarial networks (DCGANs), that have certain architectural constraints, and demonstrate that they are a strong candidate for unsupervised learning. Training on various image datasets, we show convincing evidence that our deep convolutional adversarial pair learns a hierarchy of representations from object parts to scenes in both the generator and discriminator. Additionally, we use the learned features for novel tasks - demonstrating their applicability as general image representations.\n\n[[Paper]](https://arxiv.org/abs/1511.06434) [[Code]](implementations/dcgan/dcgan.py)\n\n#### Run Example\n```\n$ cd implementations/dcgan/\n$ python3 dcgan.py\n```\n\n\n ![](\"assets/dcgan.gif\")
\n
\n\n### DiscoGAN\n_Learning to Discover Cross-Domain Relations with Generative Adversarial Networks_\n\n#### Authors\nTaeksoo Kim, Moonsu Cha, Hyunsoo Kim, Jung Kwon Lee, Jiwon Kim\n\n#### Abstract\nWhile humans easily recognize relations between data from different domains without any supervision, learning to automatically discover them is in general very challenging and needs many ground-truth pairs that illustrate the relations. To avoid costly pairing, we address the task of discovering cross-domain relations given unpaired data. We propose a method based on generative adversarial networks that learns to discover relations between different domains (DiscoGAN). Using the discovered relations, our proposed network successfully transfers style from one domain to another while preserving key attributes such as orientation and face identity.\n\n[[Paper]](https://arxiv.org/abs/1703.05192) [[Code]](implementations/discogan/discogan.py)\n\n\n ![](\"http://eriklindernoren.se/images/discogan_architecture.png\")
\n
\n\n#### Run Example\n```\n$ cd data/\n$ bash download_pix2pix_dataset.sh edges2shoes\n$ cd ../implementations/discogan/\n$ python3 discogan.py --dataset_name edges2shoes\n```\n\n\n ![](\"assets/discogan.png\")
\n
\n\n Rows from top to bottom: (1) Real image from domain A (2) Translated image from
\n domain A (3) Reconstructed image from domain A (4) Real image from domain B (5)
\n Translated image from domain B (6) Reconstructed image from domain B\n
\n\n### DRAGAN\n_On Convergence and Stability of GANs_\n\n#### Authors\nNaveen Kodali, Jacob Abernethy, James Hays, Zsolt Kira\n\n#### Abstract\nWe propose studying GAN training dynamics as regret minimization, which is in contrast to the popular view that there is consistent minimization of a divergence between real and generated distributions. We analyze the convergence of GAN training from this new point of view to understand why mode collapse happens. We hypothesize the existence of undesirable local equilibria in this non-convex game to be responsible for mode collapse. We observe that these local equilibria often exhibit sharp gradients of the discriminator function around some real data points. We demonstrate that these degenerate local equilibria can be avoided with a gradient penalty scheme called DRAGAN. We show that DRAGAN enables faster training, achieves improved stability with fewer mode collapses, and leads to generator networks with better modeling performance across a variety of architectures and objective functions.\n\n[[Paper]](https://arxiv.org/abs/1705.07215) [[Code]](implementations/dragan/dragan.py)\n\n#### Run Example\n```\n$ cd implementations/dragan/\n$ python3 dragan.py\n```\n\n### DualGAN\n_DualGAN: Unsupervised Dual Learning for Image-to-Image Translation_\n\n#### Authors\nZili Yi, Hao Zhang, Ping Tan, Minglun Gong\n\n#### Abstract\nConditional Generative Adversarial Networks (GANs) for cross-domain image-to-image translation have made much progress recently. Depending on the task complexity, thousands to millions of labeled image pairs are needed to train a conditional GAN. However, human labeling is expensive, even impractical, and large quantities of data may not always be available. Inspired by dual learning from natural language translation, we develop a novel dual-GAN mechanism, which enables image translators to be trained from two sets of unlabeled images from two domains. In our architecture, the primal GAN learns to translate images from domain U to those in domain V, while the dual GAN learns to invert the task. The closed loop made by the primal and dual tasks allows images from either domain to be translated and then reconstructed. Hence a loss function that accounts for the reconstruction error of images can be used to train the translators. Experiments on multiple image translation tasks with unlabeled data show considerable performance gain of DualGAN over a single GAN. For some tasks, DualGAN can even achieve comparable or slightly better results than conditional GAN trained on fully labeled data.\n\n[[Paper]](https://arxiv.org/abs/1704.02510) [[Code]](implementations/dualgan/dualgan.py)\n\n\n#### Run Example\n```\n$ cd data/\n$ bash download_pix2pix_dataset.sh facades\n$ cd ../implementations/dualgan/\n$ python3 dualgan.py --dataset_name facades\n```\n\n### Energy-Based GAN\n_Energy-based Generative Adversarial Network_\n\n#### Authors\nJunbo Zhao, Michael Mathieu, Yann LeCun\n\n#### Abstract\nWe introduce the \"Energy-based Generative Adversarial Network\" model (EBGAN) which views the discriminator as an energy function that attributes low energies to the regions near the data manifold and higher energies to other regions. Similar to the probabilistic GANs, a generator is seen as being trained to produce contrastive samples with minimal energies, while the discriminator is trained to assign high energies to these generated samples. Viewing the discriminator as an energy function allows to use a wide variety of architectures and loss functionals in addition to the usual binary classifier with logistic output. Among them, we show one instantiation of EBGAN framework as using an auto-encoder architecture, with the energy being the reconstruction error, in place of the discriminator. We show that this form of EBGAN exhibits more stable behavior than regular GANs during training. We also show that a single-scale architecture can be trained to generate high-resolution images.\n\n[[Paper]](https://arxiv.org/abs/1609.03126) [[Code]](implementations/ebgan/ebgan.py)\n\n#### Run Example\n```\n$ cd implementations/ebgan/\n$ python3 ebgan.py\n```\n\n### Enhanced Super-Resolution GAN\n_ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks_\n\n#### Authors\nXintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Chen Change Loy, Yu Qiao, Xiaoou Tang\n\n#### Abstract\nThe Super-Resolution Generative Adversarial Network (SRGAN) is a seminal work that is capable of generating realistic textures during single image super-resolution. However, the hallucinated details are often accompanied with unpleasant artifacts. To further enhance the visual quality, we thoroughly study three key components of SRGAN - network architecture, adversarial loss and perceptual loss, and improve each of them to derive an Enhanced SRGAN (ESRGAN). In particular, we introduce the Residual-in-Residual Dense Block (RRDB) without batch normalization as the basic network building unit. Moreover, we borrow the idea from relativistic GAN to let the discriminator predict relative realness instead of the absolute value. Finally, we improve the perceptual loss by using the features before activation, which could provide stronger supervision for brightness consistency and texture recovery. Benefiting from these improvements, the proposed ESRGAN achieves consistently better visual quality with more realistic and natural textures than SRGAN and won the first place in the PIRM2018-SR Challenge. The code is available at [this https URL](https://github.com/xinntao/ESRGAN).\n\n[[Paper]](https://arxiv.org/abs/1809.00219) [[Code]](implementations/esrgan/esrgan.py)\n\n\n#### Run Example\n```\n$ cd implementations/esrgan/\n\n$ python3 esrgan.py\n```\n\n\n ![](\"assets/enhanced_superresgan.png\")
\n
\n\n Nearest Neighbor Upsampling | ESRGAN\n
\n\n### GAN\n_Generative Adversarial Network_\n\n#### Authors\nIan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, Yoshua Bengio\n\n#### Abstract\nWe propose a new framework for estimating generative models via an adversarial process, in which we simultaneously train two models: a generative model G that captures the data distribution, and a discriminative model D that estimates the probability that a sample came from the training data rather than G. The training procedure for G is to maximize the probability of D making a mistake. This framework corresponds to a minimax two-player game. In the space of arbitrary functions G and D, a unique solution exists, with G recovering the training data distribution and D equal to 1/2 everywhere. In the case where G and D are defined by multilayer perceptrons, the entire system can be trained with backpropagation. There is no need for any Markov chains or unrolled approximate inference networks during either training or generation of samples. Experiments demonstrate the potential of the framework through qualitative and quantitative evaluation of the generated samples.\n\n[[Paper]](https://arxiv.org/abs/1406.2661) [[Code]](implementations/gan/gan.py)\n\n#### Run Example\n```\n$ cd implementations/gan/\n$ python3 gan.py\n```\n\n\n ![](\"assets/gan.gif\")
\n
\n\n### InfoGAN\n_InfoGAN: Interpretable Representation Learning by Information Maximizing Generative Adversarial Nets_\n\n#### Authors\nXi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, Pieter Abbeel\n\n#### Abstract\nThis paper describes InfoGAN, an information-theoretic extension to the Generative Adversarial Network that is able to learn disentangled representations in a completely unsupervised manner. InfoGAN is a generative adversarial network that also maximizes the mutual information between a small subset of the latent variables and the observation. We derive a lower bound to the mutual information objective that can be optimized efficiently, and show that our training procedure can be interpreted as a variation of the Wake-Sleep algorithm. Specifically, InfoGAN successfully disentangles writing styles from digit shapes on the MNIST dataset, pose from lighting of 3D rendered images, and background digits from the central digit on the SVHN dataset. It also discovers visual concepts that include hair styles, presence/absence of eyeglasses, and emotions on the CelebA face dataset. Experiments show that InfoGAN learns interpretable representations that are competitive with representations learned by existing fully supervised methods.\n\n[[Paper]](https://arxiv.org/abs/1606.03657) [[Code]](implementations/infogan/infogan.py)\n\n#### Run Example\n```\n$ cd implementations/infogan/\n$ python3 infogan.py\n```\n\n\n ![](\"assets/infogan.gif\")
\n
\n\n Result of varying categorical latent variable by column.\n
\n\n\n ![](\"assets/infogan.png\")
\n
\n\n Result of varying continuous latent variable by row.\n
\n\n### Least Squares GAN\n_Least Squares Generative Adversarial Networks_\n\n#### Authors\nXudong Mao, Qing Li, Haoran Xie, Raymond Y.K. Lau, Zhen Wang, Stephen Paul Smolley\n\n#### Abstract\nUnsupervised learning with generative adversarial networks (GANs) has proven hugely successful. Regular GANs hypothesize the discriminator as a classifier with the sigmoid cross entropy loss function. However, we found that this loss function may lead to the vanishing gradients problem during the learning process. To overcome such a problem, we propose in this paper the Least Squares Generative Adversarial Networks (LSGANs) which adopt the least squares loss function for the discriminator. We show that minimizing the objective function of LSGAN yields minimizing the Pearson χ2 divergence. There are two benefits of LSGANs over regular GANs. First, LSGANs are able to generate higher quality images than regular GANs. Second, LSGANs perform more stable during the learning process. We evaluate LSGANs on five scene datasets and the experimental results show that the images generated by LSGANs are of better quality than the ones generated by regular GANs. We also conduct two comparison experiments between LSGANs and regular GANs to illustrate the stability of LSGANs.\n\n[[Paper]](https://arxiv.org/abs/1611.04076) [[Code]](implementations/lsgan/lsgan.py)\n\n#### Run Example\n```\n$ cd implementations/lsgan/\n$ python3 lsgan.py\n```\n\n\n### MUNIT\n_Multimodal Unsupervised Image-to-Image Translation_\n\n#### Authors\nXun Huang, Ming-Yu Liu, Serge Belongie, Jan Kautz\n\n#### Abstract\nUnsupervised image-to-image translation is an important and challenging problem in computer vision. Given an image in the source domain, the goal is to learn the conditional distribution of corresponding images in the target domain, without seeing any pairs of corresponding images. While this conditional distribution is inherently multimodal, existing approaches make an overly simplified assumption, modeling it as a deterministic one-to-one mapping. As a result, they fail to generate diverse outputs from a given source domain image. To address this limitation, we propose a Multimodal Unsupervised Image-to-image Translation (MUNIT) framework. We assume that the image representation can be decomposed into a content code that is domain-invariant, and a style code that captures domain-specific properties. To translate an image to another domain, we recombine its content code with a random style code sampled from the style space of the target domain. We analyze the proposed framework and establish several theoretical results. Extensive experiments with comparisons to the state-of-the-art approaches further demonstrates the advantage of the proposed framework. Moreover, our framework allows users to control the style of translation outputs by providing an example style image. Code and pretrained models are available at [this https URL](https://github.com/nvlabs/MUNIT)\n\n[[Paper]](https://arxiv.org/abs/1804.04732) [[Code]](implementations/munit/munit.py)\n\n#### Run Example\n```\n$ cd data/\n$ bash download_pix2pix_dataset.sh edges2shoes\n$ cd ../implementations/munit/\n$ python3 munit.py --dataset_name edges2shoes\n```\n\n\n ![](\"assets/munit.png\")
\n
\n\n Results by varying the style code.\n
\n\n### Pix2Pix\n_Unpaired Image-to-Image Translation with Conditional Adversarial Networks_\n\n#### Authors\nPhillip Isola, Jun-Yan Zhu, Tinghui Zhou, Alexei A. Efros\n\n#### Abstract\nWe investigate conditional adversarial networks as a general-purpose solution to image-to-image translation problems. These networks not only learn the mapping from input image to output image, but also learn a loss function to train this mapping. This makes it possible to apply the same generic approach to problems that traditionally would require very different loss formulations. We demonstrate that this approach is effective at synthesizing photos from label maps, reconstructing objects from edge maps, and colorizing images, among other tasks. Indeed, since the release of the pix2pix software associated with this paper, a large number of internet users (many of them artists) have posted their own experiments with our system, further demonstrating its wide applicability and ease of adoption without the need for parameter tweaking. As a community, we no longer hand-engineer our mapping functions, and this work suggests we can achieve reasonable results without hand-engineering our loss functions either.\n\n[[Paper]](https://arxiv.org/abs/1611.07004) [[Code]](implementations/pix2pix/pix2pix.py)\n\n\n ![](\"http://eriklindernoren.se/images/pix2pix_architecture.png\")
\n
\n\n#### Run Example\n```\n$ cd data/\n$ bash download_pix2pix_dataset.sh facades\n$ cd ../implementations/pix2pix/\n$ python3 pix2pix.py --dataset_name facades\n```\n\n\n ![](\"assets/pix2pix.png\")
\n
\n\n Rows from top to bottom: (1) The condition for the generator (2) Generated image
\n based of condition (3) The true corresponding image to the condition\n
\n\n### PixelDA\n_Unsupervised Pixel-Level Domain Adaptation with Generative Adversarial Networks_\n\n#### Authors\nKonstantinos Bousmalis, Nathan Silberman, David Dohan, Dumitru Erhan, Dilip Krishnan\n\n#### Abstract\nCollecting well-annotated image datasets to train modern machine learning algorithms is prohibitively expensive for many tasks. One appealing alternative is rendering synthetic data where ground-truth annotations are generated automatically. Unfortunately, models trained purely on rendered images often fail to generalize to real images. To address this shortcoming, prior work introduced unsupervised domain adaptation algorithms that attempt to map representations between the two domains or learn to extract features that are domain-invariant. In this work, we present a new approach that learns, in an unsupervised manner, a transformation in the pixel space from one domain to the other. Our generative adversarial network (GAN)-based method adapts source-domain images to appear as if drawn from the target domain. Our approach not only produces plausible samples, but also outperforms the state-of-the-art on a number of unsupervised domain adaptation scenarios by large margins. Finally, we demonstrate that the adaptation process generalizes to object classes unseen during training.\n\n[[Paper]](https://arxiv.org/abs/1612.05424) [[Code]](implementations/pixelda/pixelda.py)\n\n#### MNIST to MNIST-M Classification\nTrains a classifier on images that have been translated from the source domain (MNIST) to the target domain (MNIST-M) using the annotations of the source domain images. The classification network is trained jointly with the generator network to optimize the generator for both providing a proper domain translation and also for preserving the semantics of the source domain image. The classification network trained on translated images is compared to the naive solution of training a classifier on MNIST and evaluating it on MNIST-M. The naive model manages a 55% classification accuracy on MNIST-M while the one trained during domain adaptation achieves a 95% classification accuracy.\n\n```\n$ cd implementations/pixelda/\n$ python3 pixelda.py\n``` \n| Method | Accuracy |\n| ------------ |:---------:|\n| Naive | 55% |\n| PixelDA | 95% |\n\n\n ![](\"assets/pixelda.png\")
\n
\n\n Rows from top to bottom: (1) Real images from MNIST (2) Translated images from
\n MNIST to MNIST-M (3) Examples of images from MNIST-M\n
\n\n### Relativistic GAN\n_The relativistic discriminator: a key element missing from standard GAN_\n\n#### Authors\nAlexia Jolicoeur-Martineau\n\n#### Abstract\nIn standard generative adversarial network (SGAN), the discriminator estimates the probability that the input data is real. The generator is trained to increase the probability that fake data is real. We argue that it should also simultaneously decrease the probability that real data is real because 1) this would account for a priori knowledge that half of the data in the mini-batch is fake, 2) this would be observed with divergence minimization, and 3) in optimal settings, SGAN would be equivalent to integral probability metric (IPM) GANs.\nWe show that this property can be induced by using a relativistic discriminator which estimate the probability that the given real data is more realistic than a randomly sampled fake data. We also present a variant in which the discriminator estimate the probability that the given real data is more realistic than fake data, on average. We generalize both approaches to non-standard GAN loss functions and we refer to them respectively as Relativistic GANs (RGANs) and Relativistic average GANs (RaGANs). We show that IPM-based GANs are a subset of RGANs which use the identity function.\nEmpirically, we observe that 1) RGANs and RaGANs are significantly more stable and generate higher quality data samples than their non-relativistic counterparts, 2) Standard RaGAN with gradient penalty generate data of better quality than WGAN-GP while only requiring a single discriminator update per generator update (reducing the time taken for reaching the state-of-the-art by 400%), and 3) RaGANs are able to generate plausible high resolutions images (256x256) from a very small sample (N=2011), while GAN and LSGAN cannot; these images are of significantly better quality than the ones generated by WGAN-GP and SGAN with spectral normalization.\n\n\n[[Paper]](https://arxiv.org/abs/1807.00734) [[Code]](implementations/relativistic_gan/relativistic_gan.py)\n\n#### Run Example\n```\n$ cd implementations/relativistic_gan/\n$ python3 relativistic_gan.py # Relativistic Standard GAN\n$ python3 relativistic_gan.py --rel_avg_gan # Relativistic Average GAN\n```\n\n### Semi-Supervised GAN\n_Semi-Supervised Generative Adversarial Network_\n\n#### Authors\nAugustus Odena\n\n#### Abstract\nWe extend Generative Adversarial Networks (GANs) to the semi-supervised context by forcing the discriminator network to output class labels. We train a generative model G and a discriminator D on a dataset with inputs belonging to one of N classes. At training time, D is made to predict which of N+1 classes the input belongs to, where an extra class is added to correspond to the outputs of G. We show that this method can be used to create a more data-efficient classifier and that it allows for generating higher quality samples than a regular GAN.\n\n[[Paper]](https://arxiv.org/abs/1606.01583) [[Code]](implementations/sgan/sgan.py)\n\n#### Run Example\n```\n$ cd implementations/sgan/\n$ python3 sgan.py\n```\n\n### Softmax GAN\n_Softmax GAN_\n\n#### Authors\nMin Lin\n\n#### Abstract\nSoftmax GAN is a novel variant of Generative Adversarial Network (GAN). The key idea of Softmax GAN is to replace the classification loss in the original GAN with a softmax cross-entropy loss in the sample space of one single batch. In the adversarial learning of N real training samples and M generated samples, the target of discriminator training is to distribute all the probability mass to the real samples, each with probability 1M, and distribute zero probability to generated data. In the generator training phase, the target is to assign equal probability to all data points in the batch, each with probability 1M+N. While the original GAN is closely related to Noise Contrastive Estimation (NCE), we show that Softmax GAN is the Importance Sampling version of GAN. We futher demonstrate with experiments that this simple change stabilizes GAN training.\n\n[[Paper]](https://arxiv.org/abs/1704.06191) [[Code]](implementations/softmax_gan/softmax_gan.py)\n\n#### Run Example\n```\n$ cd implementations/softmax_gan/\n$ python3 softmax_gan.py\n```\n\n### StarGAN\n_StarGAN: Unified Generative Adversarial Networks for Multi-Domain Image-to-Image Translation_\n\n#### Authors\nYunjey Choi, Minje Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, Jaegul Choo\n\n#### Abstract\nRecent studies have shown remarkable success in image-to-image translation for two domains. However, existing approaches have limited scalability and robustness in handling more than two domains, since different models should be built independently for every pair of image domains. To address this limitation, we propose StarGAN, a novel and scalable approach that can perform image-to-image translations for multiple domains using only a single model. Such a unified model architecture of StarGAN allows simultaneous training of multiple datasets with different domains within a single network. This leads to StarGAN's superior quality of translated images compared to existing models as well as the novel capability of flexibly translating an input image to any desired target domain. We empirically demonstrate the effectiveness of our approach on a facial attribute transfer and a facial expression synthesis tasks.\n\n[[Paper]](https://arxiv.org/abs/1711.09020) [[Code]](implementations/stargan/stargan.py)\n\n#### Run Example\n```\n$ cd implementations/stargan/\n\n$ python3 stargan.py\n```\n\n\n ![](\"assets/stargan.png\")
\n
\n\n Original | Black Hair | Blonde Hair | Brown Hair | Gender Flip | Aged\n
\n\n### Super-Resolution GAN\n_Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network_\n\n#### Authors\nChristian Ledig, Lucas Theis, Ferenc Huszar, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, Wenzhe Shi\n\n#### Abstract\nDespite the breakthroughs in accuracy and speed of single image super-resolution using faster and deeper convolutional neural networks, one central problem remains largely unsolved: how do we recover the finer texture details when we super-resolve at large upscaling factors? The behavior of optimization-based super-resolution methods is principally driven by the choice of the objective function. Recent work has largely focused on minimizing the mean squared reconstruction error. The resulting estimates have high peak signal-to-noise ratios, but they are often lacking high-frequency details and are perceptually unsatisfying in the sense that they fail to match the fidelity expected at the higher resolution. In this paper, we present SRGAN, a generative adversarial network (GAN) for image super-resolution (SR). To our knowledge, it is the first framework capable of inferring photo-realistic natural images for 4x upscaling factors. To achieve this, we propose a perceptual loss function which consists of an adversarial loss and a content loss. The adversarial loss pushes our solution to the natural image manifold using a discriminator network that is trained to differentiate between the super-resolved images and original photo-realistic images. In addition, we use a content loss motivated by perceptual similarity instead of similarity in pixel space. Our deep residual network is able to recover photo-realistic textures from heavily downsampled images on public benchmarks. An extensive mean-opinion-score (MOS) test shows hugely significant gains in perceptual quality using SRGAN. The MOS scores obtained with SRGAN are closer to those of the original high-resolution images than to those obtained with any state-of-the-art method.\n\n[[Paper]](https://arxiv.org/abs/1609.04802) [[Code]](implementations/srgan/srgan.py)\n\n\n ![](\"http://eriklindernoren.se/images/superresgan.png\")
\n
\n\n#### Run Example\n```\n$ cd implementations/srgan/\n\n$ python3 srgan.py\n```\n\n\n ![](\"assets/superresgan.png\")
\n
\n\n Nearest Neighbor Upsampling | SRGAN\n
\n\n### UNIT\n_Unsupervised Image-to-Image Translation Networks_\n\n#### Authors\nMing-Yu Liu, Thomas Breuel, Jan Kautz\n\n#### Abstract\nUnsupervised image-to-image translation aims at learning a joint distribution of images in different domains by using images from the marginal distributions in individual domains. Since there exists an infinite set of joint distributions that can arrive the given marginal distributions, one could infer nothing about the joint distribution from the marginal distributions without additional assumptions. To address the problem, we make a shared-latent space assumption and propose an unsupervised image-to-image translation framework based on Coupled GANs. We compare the proposed framework with competing approaches and present high quality image translation results on various challenging unsupervised image translation tasks, including street scene image translation, animal image translation, and face image translation. We also apply the proposed framework to domain adaptation and achieve state-of-the-art performance on benchmark datasets. Code and additional results are available in this [https URL](https://github.com/mingyuliutw/unit).\n\n[[Paper]](https://arxiv.org/abs/1703.00848) [[Code]](implementations/unit/unit.py)\n\n#### Run Example\n```\n$ cd data/\n$ bash download_cyclegan_dataset.sh apple2orange\n$ cd implementations/unit/\n$ python3 unit.py --dataset_name apple2orange\n```\n\n### Wasserstein GAN\n_Wasserstein GAN_\n\n#### Authors\nMartin Arjovsky, Soumith Chintala, Léon Bottou\n\n#### Abstract\nWe introduce a new algorithm named WGAN, an alternative to traditional GAN training. In this new model, we show that we can improve the stability of learning, get rid of problems like mode collapse, and provide meaningful learning curves useful for debugging and hyperparameter searches. Furthermore, we show that the corresponding optimization problem is sound, and provide extensive theoretical work highlighting the deep connections to other distances between distributions.\n\n[[Paper]](https://arxiv.org/abs/1701.07875) [[Code]](implementations/wgan/wgan.py)\n\n#### Run Example\n```\n$ cd implementations/wgan/\n$ python3 wgan.py\n```\n\n### Wasserstein GAN GP\n_Improved Training of Wasserstein GANs_\n\n#### Authors\nIshaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, Aaron Courville\n\n#### Abstract\nGenerative Adversarial Networks (GANs) are powerful generative models, but suffer from training instability. The recently proposed Wasserstein GAN (WGAN) makes progress toward stable training of GANs, but sometimes can still generate only low-quality samples or fail to converge. We find that these problems are often due to the use of weight clipping in WGAN to enforce a Lipschitz constraint on the critic, which can lead to undesired behavior. We propose an alternative to clipping weights: penalize the norm of gradient of the critic with respect to its input. Our proposed method performs better than standard WGAN and enables stable training of a wide variety of GAN architectures with almost no hyperparameter tuning, including 101-layer ResNets and language models over discrete data. We also achieve high quality generations on CIFAR-10 and LSUN bedrooms.\n\n[[Paper]](https://arxiv.org/abs/1704.00028) [[Code]](implementations/wgan_gp/wgan_gp.py)\n\n#### Run Example\n```\n$ cd implementations/wgan_gp/\n$ python3 wgan_gp.py\n```\n\n\n ![](\"assets/wgan_gp.gif\")
\n
\n\n### Wasserstein GAN DIV\n_Wasserstein Divergence for GANs_\n\n#### Authors\nJiqing Wu, Zhiwu Huang, Janine Thoma, Dinesh Acharya, Luc Van Gool\n\n#### Abstract\nIn many domains of computer vision, generative adversarial networks (GANs) have achieved great success, among which the fam-\nily of Wasserstein GANs (WGANs) is considered to be state-of-the-art due to the theoretical contributions and competitive qualitative performance. However, it is very challenging to approximate the k-Lipschitz constraint required by the Wasserstein-1 metric (W-met). In this paper, we propose a novel Wasserstein divergence (W-div), which is a relaxed version of W-met and does not require the k-Lipschitz constraint.As a concrete application, we introduce a Wasserstein divergence objective for GANs (WGAN-div), which can faithfully approximate W-div through optimization. Under various settings, including progressive growing training, we demonstrate the stability of the proposed WGAN-div owing to its theoretical and practical advantages over WGANs. Also, we study the quantitative and visual performance of WGAN-div on standard image synthesis benchmarks, showing the superior performance of WGAN-div compared to the state-of-the-art methods.\n\n[[Paper]](https://arxiv.org/abs/1712.01026) [[Code]](implementations/wgan_div/wgan_div.py)\n\n#### Run Example\n```\n$ cd implementations/wgan_div/\n$ python3 wgan_div.py\n```\n\n\n ![](\"assets/wgan_div.png\")
\n
\n"
},
{
"path": "data/download_cyclegan_dataset.sh",
"content": "#!/bin/bash\n\nFILE=$1\n\nif [[ $FILE != \"ae_photos\" && $FILE != \"apple2orange\" && $FILE != \"summer2winter_yosemite\" && $FILE != \"horse2zebra\" && $FILE != \"monet2photo\" && $FILE != \"cezanne2photo\" && $FILE != \"ukiyoe2photo\" && $FILE != \"vangogh2photo\" && $FILE != \"maps\" && $FILE != \"cityscapes\" && $FILE != \"facades\" && $FILE != \"iphone2dslr_flower\" && $FILE != \"ae_photos\" ]]; then\n echo \"Available datasets are: apple2orange, summer2winter_yosemite, horse2zebra, monet2photo, cezanne2photo, ukiyoe2photo, vangogh2photo, maps, cityscapes, facades, iphone2dslr_flower, ae_photos\"\n exit 1\nfi\n\nURL=https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/$FILE.zip\nZIP_FILE=./$FILE.zip\nTARGET_DIR=./$FILE\nwget -N $URL -O $ZIP_FILE\nunzip $ZIP_FILE -d .\nrm $ZIP_FILE\n\n# Adapt to project expected directory heriarchy\nmkdir -p \"$TARGET_DIR/train\" \"$TARGET_DIR/test\"\nmv \"$TARGET_DIR/trainA\" \"$TARGET_DIR/train/A\"\nmv \"$TARGET_DIR/trainB\" \"$TARGET_DIR/train/B\"\nmv \"$TARGET_DIR/testA\" \"$TARGET_DIR/test/A\"\nmv \"$TARGET_DIR/testB\" \"$TARGET_DIR/test/B\"\n"
},
{
"path": "data/download_pix2pix_dataset.sh",
"content": "FILE=$1\nURL=https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/$FILE.tar.gz\nTAR_FILE=./$FILE.tar.gz\nTARGET_DIR=./$FILE/\nwget -N $URL -O $TAR_FILE\nmkdir $TARGET_DIR\ntar -zxvf $TAR_FILE -C ./\nrm $TAR_FILE\n"
},
{
"path": "implementations/aae/aae.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=10, help=\"dimensionality of the latent code\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef reparameterization(mu, logvar):\n std = torch.exp(logvar / 2)\n sampled_z = Variable(Tensor(np.random.normal(0, 1, (mu.size(0), opt.latent_dim))))\n z = sampled_z * std + mu\n return z\n\n\nclass Encoder(nn.Module):\n def __init__(self):\n super(Encoder, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 512),\n nn.BatchNorm1d(512),\n nn.LeakyReLU(0.2, inplace=True),\n )\n\n self.mu = nn.Linear(512, opt.latent_dim)\n self.logvar = nn.Linear(512, opt.latent_dim)\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n x = self.model(img_flat)\n mu = self.mu(x)\n logvar = self.logvar(x)\n z = reparameterization(mu, logvar)\n return z\n\n\nclass Decoder(nn.Module):\n def __init__(self):\n super(Decoder, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(opt.latent_dim, 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 512),\n nn.BatchNorm1d(512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, int(np.prod(img_shape))),\n nn.Tanh(),\n )\n\n def forward(self, z):\n img_flat = self.model(z)\n img = img_flat.view(img_flat.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(opt.latent_dim, 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n nn.Sigmoid(),\n )\n\n def forward(self, z):\n validity = self.model(z)\n return validity\n\n\n# Use binary cross-entropy loss\nadversarial_loss = torch.nn.BCELoss()\npixelwise_loss = torch.nn.L1Loss()\n\n# Initialize generator and discriminator\nencoder = Encoder()\ndecoder = Decoder()\ndiscriminator = Discriminator()\n\nif cuda:\n encoder.cuda()\n decoder.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n pixelwise_loss.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(encoder.parameters(), decoder.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef sample_image(n_row, batches_done):\n \"\"\"Saves a grid of generated digits\"\"\"\n # Sample noise\n z = Variable(Tensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))\n gen_imgs = decoder(z)\n save_image(gen_imgs.data, \"images/%d.png\" % batches_done, nrow=n_row, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n encoded_imgs = encoder(real_imgs)\n decoded_imgs = decoder(encoded_imgs)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = 0.001 * adversarial_loss(discriminator(encoded_imgs), valid) + 0.999 * pixelwise_loss(\n decoded_imgs, real_imgs\n )\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Sample noise as discriminator ground truth\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(z), valid)\n fake_loss = adversarial_loss(discriminator(encoded_imgs.detach()), fake)\n d_loss = 0.5 * (real_loss + fake_loss)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n sample_image(n_row=10, batches_done=batches_done)\n"
},
{
"path": "implementations/acgan/acgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--n_classes\", type=int, default=10, help=\"number of classes for dataset\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.label_emb = nn.Embedding(opt.n_classes, opt.latent_dim)\n\n self.init_size = opt.img_size // 4 # Initial size before upsampling\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise, labels):\n gen_input = torch.mul(self.label_emb(labels), noise)\n out = self.l1(gen_input)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n \"\"\"Returns layers of each discriminator block\"\"\"\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.conv_blocks = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n\n # Output layers\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())\n self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.n_classes), nn.Softmax())\n\n def forward(self, img):\n out = self.conv_blocks(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n label = self.aux_layer(out)\n\n return validity, label\n\n\n# Loss functions\nadversarial_loss = torch.nn.BCELoss()\nauxiliary_loss = torch.nn.CrossEntropyLoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n auxiliary_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n\ndef sample_image(n_row, batches_done):\n \"\"\"Saves a grid of generated digits ranging from 0 to n_classes\"\"\"\n # Sample noise\n z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))\n # Get labels ranging from 0 to n_classes for n rows\n labels = np.array([num for _ in range(n_row) for num in range(n_row)])\n labels = Variable(LongTensor(labels))\n gen_imgs = generator(z, labels)\n save_image(gen_imgs.data, \"images/%d.png\" % batches_done, nrow=n_row, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, labels) in enumerate(dataloader):\n\n batch_size = imgs.shape[0]\n\n # Adversarial ground truths\n valid = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)\n fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(FloatTensor))\n labels = Variable(labels.type(LongTensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise and labels as generator input\n z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n gen_labels = Variable(LongTensor(np.random.randint(0, opt.n_classes, batch_size)))\n\n # Generate a batch of images\n gen_imgs = generator(z, gen_labels)\n\n # Loss measures generator's ability to fool the discriminator\n validity, pred_label = discriminator(gen_imgs)\n g_loss = 0.5 * (adversarial_loss(validity, valid) + auxiliary_loss(pred_label, gen_labels))\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Loss for real images\n real_pred, real_aux = discriminator(real_imgs)\n d_real_loss = (adversarial_loss(real_pred, valid) + auxiliary_loss(real_aux, labels)) / 2\n\n # Loss for fake images\n fake_pred, fake_aux = discriminator(gen_imgs.detach())\n d_fake_loss = (adversarial_loss(fake_pred, fake) + auxiliary_loss(fake_aux, gen_labels)) / 2\n\n # Total discriminator loss\n d_loss = (d_real_loss + d_fake_loss) / 2\n\n # Calculate discriminator accuracy\n pred = np.concatenate([real_aux.data.cpu().numpy(), fake_aux.data.cpu().numpy()], axis=0)\n gt = np.concatenate([labels.data.cpu().numpy(), gen_labels.data.cpu().numpy()], axis=0)\n d_acc = np.mean(np.argmax(pred, axis=1) == gt)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), 100 * d_acc, g_loss.item())\n )\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n sample_image(n_row=10, batches_done=batches_done)\n"
},
{
"path": "implementations/began/began.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=62, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"number of image channels\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise):\n out = self.l1(noise)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n # Upsampling\n self.down = nn.Sequential(nn.Conv2d(opt.channels, 64, 3, 2, 1), nn.ReLU())\n # Fully-connected layers\n self.down_size = opt.img_size // 2\n down_dim = 64 * (opt.img_size // 2) ** 2\n self.fc = nn.Sequential(\n nn.Linear(down_dim, 32),\n nn.BatchNorm1d(32, 0.8),\n nn.ReLU(inplace=True),\n nn.Linear(32, down_dim),\n nn.BatchNorm1d(down_dim),\n nn.ReLU(inplace=True),\n )\n # Upsampling\n self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv2d(64, opt.channels, 3, 1, 1))\n\n def forward(self, img):\n out = self.down(img)\n out = self.fc(out.view(out.size(0), -1))\n out = self.up(out.view(out.size(0), 64, self.down_size, self.down_size))\n return out\n\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\n# BEGAN hyper parameters\ngamma = 0.75\nlambda_k = 0.001\nk = 0.0\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = torch.mean(torch.abs(discriminator(gen_imgs) - gen_imgs))\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n d_real = discriminator(real_imgs)\n d_fake = discriminator(gen_imgs.detach())\n\n d_loss_real = torch.mean(torch.abs(d_real - real_imgs))\n d_loss_fake = torch.mean(torch.abs(d_fake - gen_imgs.detach()))\n d_loss = d_loss_real - k * d_loss_fake\n\n d_loss.backward()\n optimizer_D.step()\n\n # ----------------\n # Update weights\n # ----------------\n\n diff = torch.mean(gamma * d_loss_real - d_loss_fake)\n\n # Update weight term for fake samples\n k = k + lambda_k * diff.item()\n k = min(max(k, 0), 1) # Constraint to interval [0, 1]\n\n # Update convergence metric\n M = (d_loss_real + torch.abs(diff)).data[0]\n\n # --------------\n # Log Progress\n # --------------\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] -- M: %f, k: %f\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), M, k)\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/bgan/bgan.py",
"content": "# Reference: https://wiseodd.github.io/techblog/2017/03/07/boundary-seeking-gan/\n\nimport argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n nn.Sigmoid(),\n )\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n validity = self.model(img_flat)\n return validity\n\n\ndef boundary_seeking_loss(y_pred, y_true):\n \"\"\"\n Boundary seeking loss.\n Reference: https://wiseodd.github.io/techblog/2017/03/07/boundary-seeking-gan/\n \"\"\"\n return 0.5 * torch.mean((torch.log(y_pred) - torch.log(1 - y_pred)) ** 2)\n\n\ndiscriminator_loss = torch.nn.BCELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n discriminator_loss.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(mnist_loader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = boundary_seeking_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = discriminator_loss(discriminator(real_imgs), valid)\n fake_loss = discriminator_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(mnist_loader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(mnist_loader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/bicyclegan/bicyclegan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport datetime\nimport time\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"edges2shoes\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=8, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=128, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=128, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--latent_dim\", type=int, default=8, help=\"number of latent codes\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between saving generator samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nparser.add_argument(\"--lambda_pixel\", type=float, default=10, help=\"pixelwise loss weight\")\nparser.add_argument(\"--lambda_latent\", type=float, default=0.5, help=\"latent loss weight\")\nparser.add_argument(\"--lambda_kl\", type=float, default=0.01, help=\"kullback-leibler loss weight\")\nopt = parser.parse_args()\nprint(opt)\n\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\ncuda = True if torch.cuda.is_available() else False\n\ninput_shape = (opt.channels, opt.img_height, opt.img_width)\n\n# Loss functions\nmae_loss = torch.nn.L1Loss()\n\n# Initialize generator, encoder and discriminators\ngenerator = Generator(opt.latent_dim, input_shape)\nencoder = Encoder(opt.latent_dim, input_shape)\nD_VAE = MultiDiscriminator(input_shape)\nD_LR = MultiDiscriminator(input_shape)\n\nif cuda:\n generator = generator.cuda()\n encoder.cuda()\n D_VAE = D_VAE.cuda()\n D_LR = D_LR.cuda()\n mae_loss.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n generator.load_state_dict(torch.load(\"saved_models/%s/generator_%d.pth\" % (opt.dataset_name, opt.epoch)))\n encoder.load_state_dict(torch.load(\"saved_models/%s/encoder_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_VAE.load_state_dict(torch.load(\"saved_models/%s/D_VAE_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_LR.load_state_dict(torch.load(\"saved_models/%s/D_LR_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n generator.apply(weights_init_normal)\n D_VAE.apply(weights_init_normal)\n D_LR.apply(weights_init_normal)\n\n# Optimizers\noptimizer_E = torch.optim.Adam(encoder.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D_VAE = torch.optim.Adam(D_VAE.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D_LR = torch.optim.Adam(D_LR.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, input_shape),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, input_shape, mode=\"val\"),\n batch_size=8,\n shuffle=True,\n num_workers=1,\n)\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the validation set\"\"\"\n generator.eval()\n imgs = next(iter(val_dataloader))\n img_samples = None\n for img_A, img_B in zip(imgs[\"A\"], imgs[\"B\"]):\n # Repeat input image by number of desired columns\n real_A = img_A.view(1, *img_A.shape).repeat(opt.latent_dim, 1, 1, 1)\n real_A = Variable(real_A.type(Tensor))\n # Sample latent representations\n sampled_z = Variable(Tensor(np.random.normal(0, 1, (opt.latent_dim, opt.latent_dim))))\n # Generate samples\n fake_B = generator(real_A, sampled_z)\n # Concatenate samples horisontally\n fake_B = torch.cat([x for x in fake_B.data.cpu()], -1)\n img_sample = torch.cat((img_A, fake_B), -1)\n img_sample = img_sample.view(1, *img_sample.shape)\n # Concatenate with previous samples vertically\n img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)\n save_image(img_samples, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=8, normalize=True)\n generator.train()\n\n\ndef reparameterization(mu, logvar):\n std = torch.exp(logvar / 2)\n sampled_z = Variable(Tensor(np.random.normal(0, 1, (mu.size(0), opt.latent_dim))))\n z = sampled_z * std + mu\n return z\n\n\n# ----------\n# Training\n# ----------\n\n# Adversarial loss\nvalid = 1\nfake = 0\n\nprev_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Set model input\n real_A = Variable(batch[\"A\"].type(Tensor))\n real_B = Variable(batch[\"B\"].type(Tensor))\n\n # -------------------------------\n # Train Generator and Encoder\n # -------------------------------\n\n optimizer_E.zero_grad()\n optimizer_G.zero_grad()\n\n # ----------\n # cVAE-GAN\n # ----------\n\n # Produce output using encoding of B (cVAE-GAN)\n mu, logvar = encoder(real_B)\n encoded_z = reparameterization(mu, logvar)\n fake_B = generator(real_A, encoded_z)\n\n # Pixelwise loss of translated image by VAE\n loss_pixel = mae_loss(fake_B, real_B)\n # Kullback-Leibler divergence of encoded B\n loss_kl = 0.5 * torch.sum(torch.exp(logvar) + mu ** 2 - logvar - 1)\n # Adversarial loss\n loss_VAE_GAN = D_VAE.compute_loss(fake_B, valid)\n\n # ---------\n # cLR-GAN\n # ---------\n\n # Produce output using sampled z (cLR-GAN)\n sampled_z = Variable(Tensor(np.random.normal(0, 1, (real_A.size(0), opt.latent_dim))))\n _fake_B = generator(real_A, sampled_z)\n # cLR Loss: Adversarial loss\n loss_LR_GAN = D_LR.compute_loss(_fake_B, valid)\n\n # ----------------------------------\n # Total Loss (Generator + Encoder)\n # ----------------------------------\n\n loss_GE = loss_VAE_GAN + loss_LR_GAN + opt.lambda_pixel * loss_pixel + opt.lambda_kl * loss_kl\n\n loss_GE.backward(retain_graph=True)\n optimizer_E.step()\n\n # ---------------------\n # Generator Only Loss\n # ---------------------\n\n # Latent L1 loss\n _mu, _ = encoder(_fake_B)\n loss_latent = opt.lambda_latent * mae_loss(_mu, sampled_z)\n\n loss_latent.backward()\n optimizer_G.step()\n\n # ----------------------------------\n # Train Discriminator (cVAE-GAN)\n # ----------------------------------\n\n optimizer_D_VAE.zero_grad()\n\n loss_D_VAE = D_VAE.compute_loss(real_B, valid) + D_VAE.compute_loss(fake_B.detach(), fake)\n\n loss_D_VAE.backward()\n optimizer_D_VAE.step()\n\n # ---------------------------------\n # Train Discriminator (cLR-GAN)\n # ---------------------------------\n\n optimizer_D_LR.zero_grad()\n\n loss_D_LR = D_LR.compute_loss(real_B, valid) + D_LR.compute_loss(_fake_B.detach(), fake)\n\n loss_D_LR.backward()\n optimizer_D_LR.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D VAE_loss: %f, LR_loss: %f] [G loss: %f, pixel: %f, kl: %f, latent: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D_VAE.item(),\n loss_D_LR.item(),\n loss_GE.item(),\n loss_pixel.item(),\n loss_kl.item(),\n loss_latent.item(),\n time_left,\n )\n )\n\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(generator.state_dict(), \"saved_models/%s/generator_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(encoder.state_dict(), \"saved_models/%s/encoder_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_VAE.state_dict(), \"saved_models/%s/D_VAE_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_LR.state_dict(), \"saved_models/%s/D_LR_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/bicyclegan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, input_shape, mode=\"train\"):\n self.transform = transforms.Compose(\n [\n transforms.Resize(input_shape[-2:], Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n ]\n )\n\n self.files = sorted(glob.glob(os.path.join(root, mode) + \"/*.*\"))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n w, h = img.size\n img_A = img.crop((0, 0, w / 2, h))\n img_B = img.crop((w / 2, 0, w, h))\n\n if np.random.random() < 0.5:\n img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\")\n img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], \"RGB\")\n\n img_A = self.transform(img_A)\n img_B = self.transform(img_B)\n\n return {\"A\": img_A, \"B\": img_B}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/bicyclegan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nimport numpy as np\n\nfrom torchvision.models import resnet18\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n# U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n super(UNetDown, self).__init__()\n layers = [nn.Conv2d(in_size, out_size, 3, stride=2, padding=1, bias=False)]\n if normalize:\n layers.append(nn.BatchNorm2d(out_size, 0.8))\n layers.append(nn.LeakyReLU(0.2))\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass UNetUp(nn.Module):\n def __init__(self, in_size, out_size):\n super(UNetUp, self).__init__()\n self.model = nn.Sequential(\n nn.Upsample(scale_factor=2),\n nn.Conv2d(in_size, out_size, 3, stride=1, padding=1, bias=False),\n nn.BatchNorm2d(out_size, 0.8),\n nn.ReLU(inplace=True),\n )\n\n def forward(self, x, skip_input):\n x = self.model(x)\n x = torch.cat((x, skip_input), 1)\n return x\n\n\nclass Generator(nn.Module):\n def __init__(self, latent_dim, img_shape):\n super(Generator, self).__init__()\n channels, self.h, self.w = img_shape\n\n self.fc = nn.Linear(latent_dim, self.h * self.w)\n\n self.down1 = UNetDown(channels + 1, 64, normalize=False)\n self.down2 = UNetDown(64, 128)\n self.down3 = UNetDown(128, 256)\n self.down4 = UNetDown(256, 512)\n self.down5 = UNetDown(512, 512)\n self.down6 = UNetDown(512, 512)\n self.down7 = UNetDown(512, 512, normalize=False)\n self.up1 = UNetUp(512, 512)\n self.up2 = UNetUp(1024, 512)\n self.up3 = UNetUp(1024, 512)\n self.up4 = UNetUp(1024, 256)\n self.up5 = UNetUp(512, 128)\n self.up6 = UNetUp(256, 64)\n\n self.final = nn.Sequential(\n nn.Upsample(scale_factor=2), nn.Conv2d(128, channels, 3, stride=1, padding=1), nn.Tanh()\n )\n\n def forward(self, x, z):\n # Propogate noise through fc layer and reshape to img shape\n z = self.fc(z).view(z.size(0), 1, self.h, self.w)\n d1 = self.down1(torch.cat((x, z), 1))\n d2 = self.down2(d1)\n d3 = self.down3(d2)\n d4 = self.down4(d3)\n d5 = self.down5(d4)\n d6 = self.down6(d5)\n d7 = self.down7(d6)\n u1 = self.up1(d7, d6)\n u2 = self.up2(u1, d5)\n u3 = self.up3(u2, d4)\n u4 = self.up4(u3, d3)\n u5 = self.up5(u4, d2)\n u6 = self.up6(u5, d1)\n\n return self.final(u6)\n\n\n##############################\n# Encoder\n##############################\n\n\nclass Encoder(nn.Module):\n def __init__(self, latent_dim, input_shape):\n super(Encoder, self).__init__()\n resnet18_model = resnet18(pretrained=False)\n self.feature_extractor = nn.Sequential(*list(resnet18_model.children())[:-3])\n self.pooling = nn.AvgPool2d(kernel_size=8, stride=8, padding=0)\n # Output is mu and log(var) for reparameterization trick used in VAEs\n self.fc_mu = nn.Linear(256, latent_dim)\n self.fc_logvar = nn.Linear(256, latent_dim)\n\n def forward(self, img):\n out = self.feature_extractor(img)\n out = self.pooling(out)\n out = out.view(out.size(0), -1)\n mu = self.fc_mu(out)\n logvar = self.fc_logvar(out)\n return mu, logvar\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass MultiDiscriminator(nn.Module):\n def __init__(self, input_shape):\n super(MultiDiscriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, normalize=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.BatchNorm2d(out_filters, 0.8))\n layers.append(nn.LeakyReLU(0.2))\n return layers\n\n channels, _, _ = input_shape\n # Extracts discriminator models\n self.models = nn.ModuleList()\n for i in range(3):\n self.models.add_module(\n \"disc_%d\" % i,\n nn.Sequential(\n *discriminator_block(channels, 64, normalize=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n *discriminator_block(256, 512),\n nn.Conv2d(512, 1, 3, padding=1)\n ),\n )\n\n self.downsample = nn.AvgPool2d(in_channels, stride=2, padding=[1, 1], count_include_pad=False)\n\n def compute_loss(self, x, gt):\n \"\"\"Computes the MSE between model output and scalar gt\"\"\"\n loss = sum([torch.mean((out - gt) ** 2) for out in self.forward(x)])\n return loss\n\n def forward(self, x):\n outputs = []\n for m in self.models:\n outputs.append(m(x))\n x = self.downsample(x)\n return outputs\n"
},
{
"path": "implementations/ccgan/ccgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\nfrom PIL import Image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom datasets import *\nfrom models import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=8, help=\"size of the batches\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"img_align_celeba\", help=\"name of the dataset\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=128, help=\"size of each image dimension\")\nparser.add_argument(\"--mask_size\", type=int, default=32, help=\"size of random mask\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=500, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\ninput_shape = (opt.channels, opt.img_size, opt.img_size)\n\n# Loss function\nadversarial_loss = torch.nn.MSELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator(input_shape)\ndiscriminator = Discriminator(input_shape)\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Dataset loader\ntransforms_ = [\n transforms.Resize((opt.img_size, opt.img_size), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\ntransforms_lr = [\n transforms.Resize((opt.img_size // 4, opt.img_size // 4), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_x=transforms_, transforms_lr=transforms_lr),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef apply_random_mask(imgs):\n idx = np.random.randint(0, opt.img_size - opt.mask_size, (imgs.shape[0], 2))\n\n masked_imgs = imgs.clone()\n for i, (y1, x1) in enumerate(idx):\n y2, x2 = y1 + opt.mask_size, x1 + opt.mask_size\n masked_imgs[i, :, y1:y2, x1:x2] = -1\n\n return masked_imgs\n\n\ndef save_sample(saved_samples):\n # Generate inpainted image\n gen_imgs = generator(saved_samples[\"masked\"], saved_samples[\"lowres\"])\n # Save sample\n sample = torch.cat((saved_samples[\"masked\"].data, gen_imgs.data, saved_samples[\"imgs\"].data), -2)\n save_image(sample, \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n\n\nsaved_samples = {}\nfor epoch in range(opt.n_epochs):\n for i, batch in enumerate(dataloader):\n imgs = batch[\"x\"]\n imgs_lr = batch[\"x_lr\"]\n\n masked_imgs = apply_random_mask(imgs)\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], *discriminator.output_shape).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], *discriminator.output_shape).fill_(0.0), requires_grad=False)\n\n if cuda:\n imgs = imgs.type(Tensor)\n imgs_lr = imgs_lr.type(Tensor)\n masked_imgs = masked_imgs.type(Tensor)\n\n real_imgs = Variable(imgs)\n imgs_lr = Variable(imgs_lr)\n masked_imgs = Variable(masked_imgs)\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Generate a batch of images\n gen_imgs = generator(masked_imgs, imgs_lr)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(real_imgs), valid)\n fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = 0.5 * (real_loss + fake_loss)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n # Save first ten samples\n if not saved_samples:\n saved_samples[\"imgs\"] = real_imgs[:1].clone()\n saved_samples[\"masked\"] = masked_imgs[:1].clone()\n saved_samples[\"lowres\"] = imgs_lr[:1].clone()\n elif saved_samples[\"imgs\"].size(0) < 10:\n saved_samples[\"imgs\"] = torch.cat((saved_samples[\"imgs\"], real_imgs[:1]), 0)\n saved_samples[\"masked\"] = torch.cat((saved_samples[\"masked\"], masked_imgs[:1]), 0)\n saved_samples[\"lowres\"] = torch.cat((saved_samples[\"lowres\"], imgs_lr[:1]), 0)\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_sample(saved_samples)\n"
},
{
"path": "implementations/ccgan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_x=None, transforms_lr=None, mode='train'):\n self.transform_x = transforms.Compose(transforms_x)\n self.transform_lr = transforms.Compose(transforms_lr)\n\n self.files = sorted(glob.glob('%s/*.*' % root))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n\n x = self.transform_x(img)\n x_lr = self.transform_lr(img)\n\n return {'x': x, 'x_lr': x_lr}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/ccgan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n##############################\n# U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n super(UNetDown, self).__init__()\n model = [nn.Conv2d(in_size, out_size, 4, stride=2, padding=1, bias=False)]\n if normalize:\n model.append(nn.BatchNorm2d(out_size, 0.8))\n model.append(nn.LeakyReLU(0.2))\n if dropout:\n model.append(nn.Dropout(dropout))\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass UNetUp(nn.Module):\n def __init__(self, in_size, out_size, dropout=0.0):\n super(UNetUp, self).__init__()\n model = [\n nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),\n nn.BatchNorm2d(out_size, 0.8),\n nn.ReLU(inplace=True),\n ]\n if dropout:\n model.append(nn.Dropout(dropout))\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x, skip_input):\n x = self.model(x)\n out = torch.cat((x, skip_input), 1)\n return out\n\n\nclass Generator(nn.Module):\n def __init__(self, input_shape):\n super(Generator, self).__init__()\n channels, _, _ = input_shape\n self.down1 = UNetDown(channels, 64, normalize=False)\n self.down2 = UNetDown(64, 128)\n self.down3 = UNetDown(128 + channels, 256, dropout=0.5)\n self.down4 = UNetDown(256, 512, dropout=0.5)\n self.down5 = UNetDown(512, 512, dropout=0.5)\n self.down6 = UNetDown(512, 512, dropout=0.5)\n\n self.up1 = UNetUp(512, 512, dropout=0.5)\n self.up2 = UNetUp(1024, 512, dropout=0.5)\n self.up3 = UNetUp(1024, 256, dropout=0.5)\n self.up4 = UNetUp(512, 128)\n self.up5 = UNetUp(256 + channels, 64)\n\n final = [nn.Upsample(scale_factor=2), nn.Conv2d(128, channels, 3, 1, 1), nn.Tanh()]\n self.final = nn.Sequential(*final)\n\n def forward(self, x, x_lr):\n # U-Net generator with skip connections from encoder to decoder\n d1 = self.down1(x)\n d2 = self.down2(d1)\n d2 = torch.cat((d2, x_lr), 1)\n d3 = self.down3(d2)\n d4 = self.down4(d3)\n d5 = self.down5(d4)\n d6 = self.down6(d5)\n u1 = self.up1(d6, d5)\n u2 = self.up2(u1, d4)\n u3 = self.up3(u2, d3)\n u4 = self.up4(u3, d2)\n u5 = self.up5(u4, d1)\n\n return self.final(u5)\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n\n channels, height, width = input_shape\n # Calculate output of image discriminator (PatchGAN)\n patch_h, patch_w = int(height / 2 ** 3), int(width / 2 ** 3)\n self.output_shape = (1, patch_h, patch_w)\n\n def discriminator_block(in_filters, out_filters, stride, normalize):\n \"\"\"Returns layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 3, stride, 1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n layers = []\n in_filters = channels\n for out_filters, stride, normalize in [(64, 2, False), (128, 2, True), (256, 2, True), (512, 1, True)]:\n layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))\n in_filters = out_filters\n\n layers.append(nn.Conv2d(out_filters, 1, 3, 1, 1))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/cgan/cgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--n_classes\", type=int, default=10, help=\"number of classes for dataset\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.label_emb = nn.Embedding(opt.n_classes, opt.n_classes)\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim + opt.n_classes, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, noise, labels):\n # Concatenate label embedding and image to produce input\n gen_input = torch.cat((self.label_emb(labels), noise), -1)\n img = self.model(gen_input)\n img = img.view(img.size(0), *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.label_embedding = nn.Embedding(opt.n_classes, opt.n_classes)\n\n self.model = nn.Sequential(\n nn.Linear(opt.n_classes + int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 512),\n nn.Dropout(0.4),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 512),\n nn.Dropout(0.4),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 1),\n )\n\n def forward(self, img, labels):\n # Concatenate label embedding and image to produce input\n d_in = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)), -1)\n validity = self.model(d_in)\n return validity\n\n\n# Loss functions\nadversarial_loss = torch.nn.MSELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n\ndef sample_image(n_row, batches_done):\n \"\"\"Saves a grid of generated digits ranging from 0 to n_classes\"\"\"\n # Sample noise\n z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))\n # Get labels ranging from 0 to n_classes for n rows\n labels = np.array([num for _ in range(n_row) for num in range(n_row)])\n labels = Variable(LongTensor(labels))\n gen_imgs = generator(z, labels)\n save_image(gen_imgs.data, \"images/%d.png\" % batches_done, nrow=n_row, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, labels) in enumerate(dataloader):\n\n batch_size = imgs.shape[0]\n\n # Adversarial ground truths\n valid = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)\n fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(FloatTensor))\n labels = Variable(labels.type(LongTensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise and labels as generator input\n z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n gen_labels = Variable(LongTensor(np.random.randint(0, opt.n_classes, batch_size)))\n\n # Generate a batch of images\n gen_imgs = generator(z, gen_labels)\n\n # Loss measures generator's ability to fool the discriminator\n validity = discriminator(gen_imgs, gen_labels)\n g_loss = adversarial_loss(validity, valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Loss for real images\n validity_real = discriminator(real_imgs, labels)\n d_real_loss = adversarial_loss(validity_real, valid)\n\n # Loss for fake images\n validity_fake = discriminator(gen_imgs.detach(), gen_labels)\n d_fake_loss = adversarial_loss(validity_fake, fake)\n\n # Total discriminator loss\n d_loss = (d_real_loss + d_fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n sample_image(n_row=10, batches_done=batches_done)\n"
},
{
"path": "implementations/cluster_gan/clustergan.py",
"content": "from __future__ import print_function\n\ntry:\n import argparse\n import os\n import numpy as np\n\n from torch.autograd import Variable\n from torch.autograd import grad as torch_grad\n \n import torch\n import torchvision\n import torch.nn as nn\n import torch.nn.functional as F\n from torch.utils.data import DataLoader\n from torchvision import datasets\n import torchvision.transforms as transforms\n from torchvision.utils import save_image\n \n from itertools import chain as ichain\n\nexcept ImportError as e:\n print(e)\n raise ImportError\n\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser(description=\"ClusterGAN Training Script\")\nparser.add_argument(\"-n\", \"--n_epochs\", dest=\"n_epochs\", default=200, type=int, help=\"Number of epochs\")\nparser.add_argument(\"-b\", \"--batch_size\", dest=\"batch_size\", default=64, type=int, help=\"Batch size\")\nparser.add_argument(\"-i\", \"--img_size\", dest=\"img_size\", type=int, default=28, help=\"Size of image dimension\")\nparser.add_argument(\"-d\", \"--latent_dim\", dest=\"latent_dim\", default=30, type=int, help=\"Dimension of latent space\")\nparser.add_argument(\"-l\", \"--lr\", dest=\"learning_rate\", type=float, default=0.0001, help=\"Learning rate\")\nparser.add_argument(\"-c\", \"--n_critic\", dest=\"n_critic\", type=int, default=5, help=\"Number of training steps for discriminator per iter\")\nparser.add_argument(\"-w\", \"--wass_flag\", dest=\"wass_flag\", action='store_true', help=\"Flag for Wasserstein metric\")\nargs = parser.parse_args()\n\n\n# Sample a random latent space vector\ndef sample_z(shape=64, latent_dim=10, n_c=10, fix_class=-1, req_grad=False):\n\n assert (fix_class == -1 or (fix_class >= 0 and fix_class < n_c) ), \"Requested class %i outside bounds.\"%fix_class\n\n Tensor = torch.cuda.FloatTensor\n \n # Sample noise as generator input, zn\n zn = Variable(Tensor(0.75*np.random.normal(0, 1, (shape, latent_dim))), requires_grad=req_grad)\n\n ######### zc, zc_idx variables with grads, and zc to one-hot vector\n # Pure one-hot vector generation\n zc_FT = Tensor(shape, n_c).fill_(0)\n zc_idx = torch.empty(shape, dtype=torch.long)\n\n if (fix_class == -1):\n zc_idx = zc_idx.random_(n_c).cuda()\n zc_FT = zc_FT.scatter_(1, zc_idx.unsqueeze(1), 1.)\n else:\n zc_idx[:] = fix_class\n zc_FT[:, fix_class] = 1\n\n zc_idx = zc_idx.cuda()\n zc_FT = zc_FT.cuda()\n\n zc = Variable(zc_FT, requires_grad=req_grad)\n\n # Return components of latent space variable\n return zn, zc, zc_idx\n\ndef calc_gradient_penalty(netD, real_data, generated_data):\n # GP strength\n LAMBDA = 10\n\n b_size = real_data.size()[0]\n\n # Calculate interpolation\n alpha = torch.rand(b_size, 1, 1, 1)\n alpha = alpha.expand_as(real_data)\n alpha = alpha.cuda()\n \n interpolated = alpha * real_data.data + (1 - alpha) * generated_data.data\n interpolated = Variable(interpolated, requires_grad=True)\n interpolated = interpolated.cuda()\n\n # Calculate probability of interpolated examples\n prob_interpolated = netD(interpolated)\n\n # Calculate gradients of probabilities with respect to examples\n gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated,\n grad_outputs=torch.ones(prob_interpolated.size()).cuda(),\n create_graph=True, retain_graph=True)[0]\n\n # Gradients have shape (batch_size, num_channels, img_width, img_height),\n # so flatten to easily take norm per example in batch\n gradients = gradients.view(b_size, -1)\n\n # Derivatives of the gradient close to 0 can cause problems because of\n # the square root, so manually calculate norm and add epsilon\n gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)\n\n # Return gradient penalty\n return LAMBDA * ((gradients_norm - 1) ** 2).mean()\n\n\n# Weight Initializer\ndef initialize_weights(net):\n for m in net.modules():\n if isinstance(m, nn.Conv2d):\n m.weight.data.normal_(0, 0.02)\n m.bias.data.zero_()\n elif isinstance(m, nn.ConvTranspose2d):\n m.weight.data.normal_(0, 0.02)\n m.bias.data.zero_()\n elif isinstance(m, nn.Linear):\n m.weight.data.normal_(0, 0.02)\n m.bias.data.zero_()\n\n\n# Softmax function\ndef softmax(x):\n return F.softmax(x, dim=1)\n\n\nclass Reshape(nn.Module):\n \"\"\"\n Class for performing a reshape as a layer in a sequential model.\n \"\"\"\n def __init__(self, shape=[]):\n super(Reshape, self).__init__()\n self.shape = shape\n\n def forward(self, x):\n return x.view(x.size(0), *self.shape)\n \n def extra_repr(self):\n # (Optional)Set the extra information about this module. You can test\n # it by printing an object of this class.\n return 'shape={}'.format(\n self.shape\n )\n\n\nclass Generator_CNN(nn.Module):\n \"\"\"\n CNN to model the generator of a ClusterGAN\n Input is a vector from representation space of dimension z_dim\n output is a vector from image space of dimension X_dim\n \"\"\"\n # Architecture : FC1024_BR-FC7x7x128_BR-(64)4dc2s_BR-(1)4dc2s_S\n def __init__(self, latent_dim, n_c, x_shape, verbose=False):\n super(Generator_CNN, self).__init__()\n\n self.name = 'generator'\n self.latent_dim = latent_dim\n self.n_c = n_c\n self.x_shape = x_shape\n self.ishape = (128, 7, 7)\n self.iels = int(np.prod(self.ishape))\n self.verbose = verbose\n \n self.model = nn.Sequential(\n # Fully connected layers\n torch.nn.Linear(self.latent_dim + self.n_c, 1024),\n nn.BatchNorm1d(1024),\n nn.LeakyReLU(0.2, inplace=True),\n torch.nn.Linear(1024, self.iels),\n nn.BatchNorm1d(self.iels),\n nn.LeakyReLU(0.2, inplace=True),\n \n # Reshape to 128 x (7x7)\n Reshape(self.ishape),\n\n # Upconvolution layers\n nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1, bias=True),\n nn.BatchNorm2d(64),\n nn.LeakyReLU(0.2, inplace=True),\n \n nn.ConvTranspose2d(64, 1, 4, stride=2, padding=1, bias=True),\n nn.Sigmoid()\n )\n\n initialize_weights(self)\n\n if self.verbose:\n print(\"Setting up {}...\\n\".format(self.name))\n print(self.model)\n \n def forward(self, zn, zc):\n z = torch.cat((zn, zc), 1)\n x_gen = self.model(z)\n # Reshape for output\n x_gen = x_gen.view(x_gen.size(0), *self.x_shape)\n return x_gen\n\n\nclass Encoder_CNN(nn.Module):\n \"\"\"\n CNN to model the encoder of a ClusterGAN\n Input is vector X from image space if dimension X_dim\n Output is vector z from representation space of dimension z_dim\n \"\"\"\n def __init__(self, latent_dim, n_c, verbose=False):\n super(Encoder_CNN, self).__init__()\n\n self.name = 'encoder'\n self.channels = 1\n self.latent_dim = latent_dim\n self.n_c = n_c\n self.cshape = (128, 5, 5)\n self.iels = int(np.prod(self.cshape))\n self.lshape = (self.iels,)\n self.verbose = verbose\n \n self.model = nn.Sequential(\n # Convolutional layers\n nn.Conv2d(self.channels, 64, 4, stride=2, bias=True),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, 128, 4, stride=2, bias=True),\n nn.LeakyReLU(0.2, inplace=True),\n \n # Flatten\n Reshape(self.lshape),\n \n # Fully connected layers\n torch.nn.Linear(self.iels, 1024),\n nn.LeakyReLU(0.2, inplace=True),\n torch.nn.Linear(1024, latent_dim + n_c)\n )\n\n initialize_weights(self)\n \n if self.verbose:\n print(\"Setting up {}...\\n\".format(self.name))\n print(self.model)\n\n def forward(self, in_feat):\n z_img = self.model(in_feat)\n # Reshape for output\n z = z_img.view(z_img.shape[0], -1)\n # Separate continuous and one-hot components\n zn = z[:, 0:self.latent_dim]\n zc_logits = z[:, self.latent_dim:]\n # Softmax on zc component\n zc = softmax(zc_logits)\n return zn, zc, zc_logits\n\n\nclass Discriminator_CNN(nn.Module):\n \"\"\"\n CNN to model the discriminator of a ClusterGAN\n Input is tuple (X,z) of an image vector and its corresponding\n representation z vector. For example, if X comes from the dataset, corresponding\n z is Encoder(X), and if z is sampled from representation space, X is Generator(z)\n Output is a 1-dimensional value\n \"\"\" \n # Architecture : (64)4c2s-(128)4c2s_BL-FC1024_BL-FC1_S\n def __init__(self, wass_metric=False, verbose=False):\n super(Discriminator_CNN, self).__init__()\n \n self.name = 'discriminator'\n self.channels = 1\n self.cshape = (128, 5, 5)\n self.iels = int(np.prod(self.cshape))\n self.lshape = (self.iels,)\n self.wass = wass_metric\n self.verbose = verbose\n \n self.model = nn.Sequential(\n # Convolutional layers\n nn.Conv2d(self.channels, 64, 4, stride=2, bias=True),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, 128, 4, stride=2, bias=True),\n nn.LeakyReLU(0.2, inplace=True),\n \n # Flatten\n Reshape(self.lshape),\n \n # Fully connected layers\n torch.nn.Linear(self.iels, 1024),\n nn.LeakyReLU(0.2, inplace=True),\n torch.nn.Linear(1024, 1),\n )\n \n # If NOT using Wasserstein metric, final Sigmoid\n if (not self.wass):\n self.model = nn.Sequential(self.model, torch.nn.Sigmoid())\n\n initialize_weights(self)\n\n if self.verbose:\n print(\"Setting up {}...\\n\".format(self.name))\n print(self.model)\n\n def forward(self, img):\n # Get output\n validity = self.model(img)\n return validity\n\n\n\n# Training details\nn_epochs = args.n_epochs\nbatch_size = args.batch_size\ntest_batch_size = 5000\nlr = args.learning_rate\nb1 = 0.5\nb2 = 0.9\ndecay = 2.5*1e-5\nn_skip_iter = args.n_critic\n\n# Data dimensions\nimg_size = args.img_size\nchannels = 1\n\n# Latent space info\nlatent_dim = args.latent_dim\nn_c = 10\nbetan = 10\nbetac = 10\n\n# Wasserstein+GP metric flag\nwass_metric = args.wass_flag\n\nx_shape = (channels, img_size, img_size)\n\ncuda = True if torch.cuda.is_available() else False\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n# Loss function\nbce_loss = torch.nn.BCELoss()\nxe_loss = torch.nn.CrossEntropyLoss()\nmse_loss = torch.nn.MSELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator_CNN(latent_dim, n_c, x_shape)\nencoder = Encoder_CNN(latent_dim, n_c)\ndiscriminator = Discriminator_CNN(wass_metric=wass_metric)\n\nif cuda:\n generator.cuda()\n encoder.cuda()\n discriminator.cuda()\n bce_loss.cuda()\n xe_loss.cuda()\n mse_loss.cuda()\n \nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.ToTensor()]\n ),\n ),\n batch_size=batch_size,\n shuffle=True,\n)\n\n# Test data loader\ntestdata = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=False,\n download=True,\n transform=transforms.Compose(\n [transforms.ToTensor()]\n ),\n ),\n batch_size=batch_size,\n shuffle=True,\n)\ntest_imgs, test_labels = next(iter(testdata))\ntest_imgs = Variable(test_imgs.type(Tensor))\n\nge_chain = ichain(generator.parameters(),\n encoder.parameters())\n\noptimizer_GE = torch.optim.Adam(ge_chain, lr=lr, betas=(b1, b2), weight_decay=decay)\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))\n\n# ----------\n# Training\n# ----------\nge_l = []\nd_l = []\n\nc_zn = []\nc_zc = []\nc_i = []\n\n# Training loop \nprint('\\nBegin training session with %i epochs...\\n'%(n_epochs))\nfor epoch in range(n_epochs):\n for i, (imgs, itruth_label) in enumerate(dataloader):\n \n # Ensure generator/encoder are trainable\n generator.train()\n encoder.train()\n\n # Zero gradients for models\n generator.zero_grad()\n encoder.zero_grad()\n discriminator.zero_grad()\n \n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # ---------------------------\n # Train Generator + Encoder\n # ---------------------------\n \n optimizer_GE.zero_grad()\n \n # Sample random latent variables\n zn, zc, zc_idx = sample_z(shape=imgs.shape[0],\n latent_dim=latent_dim,\n n_c=n_c)\n\n # Generate a batch of images\n gen_imgs = generator(zn, zc)\n \n # Discriminator output from real and generated samples\n D_gen = discriminator(gen_imgs)\n D_real = discriminator(real_imgs)\n \n # Step for Generator & Encoder, n_skip_iter times less than for discriminator\n if (i % n_skip_iter == 0):\n # Encode the generated images\n enc_gen_zn, enc_gen_zc, enc_gen_zc_logits = encoder(gen_imgs)\n\n # Calculate losses for z_n, z_c\n zn_loss = mse_loss(enc_gen_zn, zn)\n zc_loss = xe_loss(enc_gen_zc_logits, zc_idx)\n\n # Check requested metric\n if wass_metric:\n # Wasserstein GAN loss\n ge_loss = torch.mean(D_gen) + betan * zn_loss + betac * zc_loss\n else:\n # Vanilla GAN loss\n valid = Variable(Tensor(gen_imgs.size(0), 1).fill_(1.0), requires_grad=False)\n v_loss = bce_loss(D_gen, valid)\n ge_loss = v_loss + betan * zn_loss + betac * zc_loss\n\n ge_loss.backward(retain_graph=True)\n optimizer_GE.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n if wass_metric:\n # Gradient penalty term\n grad_penalty = calc_gradient_penalty(discriminator, real_imgs, gen_imgs)\n\n # Wasserstein GAN loss w/gradient penalty\n d_loss = torch.mean(D_real) - torch.mean(D_gen) + grad_penalty\n \n else:\n # Vanilla GAN loss\n fake = Variable(Tensor(gen_imgs.size(0), 1).fill_(0.0), requires_grad=False)\n real_loss = bce_loss(D_real, valid)\n fake_loss = bce_loss(D_gen, fake)\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n\n # Save training losses\n d_l.append(d_loss.item())\n ge_l.append(ge_loss.item())\n\n\n # Generator in eval mode\n generator.eval()\n encoder.eval()\n\n # Set number of examples for cycle calcs\n n_sqrt_samp = 5\n n_samp = n_sqrt_samp * n_sqrt_samp\n\n\n ## Cycle through test real -> enc -> gen\n t_imgs, t_label = test_imgs.data, test_labels\n # Encode sample real instances\n e_tzn, e_tzc, e_tzc_logits = encoder(t_imgs)\n # Generate sample instances from encoding\n teg_imgs = generator(e_tzn, e_tzc)\n # Calculate cycle reconstruction loss\n img_mse_loss = mse_loss(t_imgs, teg_imgs)\n # Save img reco cycle loss\n c_i.append(img_mse_loss.item())\n \n\n ## Cycle through randomly sampled encoding -> generator -> encoder\n zn_samp, zc_samp, zc_samp_idx = sample_z(shape=n_samp,\n latent_dim=latent_dim,\n n_c=n_c)\n # Generate sample instances\n gen_imgs_samp = generator(zn_samp, zc_samp)\n\n # Encode sample instances\n zn_e, zc_e, zc_e_logits = encoder(gen_imgs_samp)\n\n # Calculate cycle latent losses\n lat_mse_loss = mse_loss(zn_e, zn_samp)\n lat_xe_loss = xe_loss(zc_e_logits, zc_samp_idx)\n\n # Save latent space cycle losses\n c_zn.append(lat_mse_loss.item())\n c_zc.append(lat_xe_loss.item())\n \n # Save cycled and generated examples!\n r_imgs, i_label = real_imgs.data[:n_samp], itruth_label[:n_samp]\n e_zn, e_zc, e_zc_logits = encoder(r_imgs)\n reg_imgs = generator(e_zn, e_zc)\n save_image(reg_imgs.data[:n_samp],\n 'images/cycle_reg_%06i.png' %(epoch), \n nrow=n_sqrt_samp, normalize=True)\n save_image(gen_imgs_samp.data[:n_samp],\n 'images/gen_%06i.png' %(epoch), \n nrow=n_sqrt_samp, normalize=True)\n \n ## Generate samples for specified classes\n stack_imgs = []\n for idx in range(n_c):\n # Sample specific class\n zn_samp, zc_samp, zc_samp_idx = sample_z(shape=n_c,\n latent_dim=latent_dim,\n n_c=n_c,\n fix_class=idx)\n\n # Generate sample instances\n gen_imgs_samp = generator(zn_samp, zc_samp)\n\n if (len(stack_imgs) == 0):\n stack_imgs = gen_imgs_samp\n else:\n stack_imgs = torch.cat((stack_imgs, gen_imgs_samp), 0)\n\n # Save class-specified generated examples!\n save_image(stack_imgs,\n 'images/gen_classes_%06i.png' %(epoch),\n nrow=n_c, normalize=True)\n \n\n print (\"[Epoch %d/%d] \\n\"\\\n \"\\tModel Losses: [D: %f] [GE: %f]\" % (epoch, \n n_epochs, \n d_loss.item(),\n ge_loss.item())\n )\n \n print(\"\\tCycle Losses: [x: %f] [z_n: %f] [z_c: %f]\"%(img_mse_loss.item(), \n lat_mse_loss.item(), \n lat_xe_loss.item())\n )\n"
},
{
"path": "implementations/cogan/cogan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport scipy\nimport itertools\n\nimport mnistm\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=32, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Linear\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass CoupledGenerators(nn.Module):\n def __init__(self):\n super(CoupledGenerators, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.fc = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.shared_conv = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n )\n self.G1 = nn.Sequential(\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n self.G2 = nn.Sequential(\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise):\n out = self.fc(noise)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img_emb = self.shared_conv(out)\n img1 = self.G1(img_emb)\n img2 = self.G2(img_emb)\n return img1, img2\n\n\nclass CoupledDiscriminators(nn.Module):\n def __init__(self):\n super(CoupledDiscriminators, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n block.extend([nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)])\n return block\n\n self.shared_conv = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n self.D1 = nn.Linear(128 * ds_size ** 2, 1)\n self.D2 = nn.Linear(128 * ds_size ** 2, 1)\n\n def forward(self, img1, img2):\n # Determine validity of first image\n out = self.shared_conv(img1)\n out = out.view(out.shape[0], -1)\n validity1 = self.D1(out)\n # Determine validity of second image\n out = self.shared_conv(img2)\n out = out.view(out.shape[0], -1)\n validity2 = self.D2(out)\n\n return validity1, validity2\n\n\n# Loss function\nadversarial_loss = torch.nn.MSELoss()\n\n# Initialize models\ncoupled_generators = CoupledGenerators()\ncoupled_discriminators = CoupledDiscriminators()\n\nif cuda:\n coupled_generators.cuda()\n coupled_discriminators.cuda()\n\n# Initialize weights\ncoupled_generators.apply(weights_init_normal)\ncoupled_discriminators.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader1 = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\nos.makedirs(\"../../data/mnistm\", exist_ok=True)\ndataloader2 = torch.utils.data.DataLoader(\n mnistm.MNISTM(\n \"../../data/mnistm\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [\n transforms.Resize(opt.img_size),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n ]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(coupled_generators.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(coupled_discriminators.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, ((imgs1, _), (imgs2, _)) in enumerate(zip(dataloader1, dataloader2)):\n\n batch_size = imgs1.shape[0]\n\n # Adversarial ground truths\n valid = Variable(Tensor(batch_size, 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(batch_size, 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n imgs1 = Variable(imgs1.type(Tensor).expand(imgs1.size(0), 3, opt.img_size, opt.img_size))\n imgs2 = Variable(imgs2.type(Tensor))\n\n # ------------------\n # Train Generators\n # ------------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs1, gen_imgs2 = coupled_generators(z)\n # Determine validity of generated images\n validity1, validity2 = coupled_discriminators(gen_imgs1, gen_imgs2)\n\n g_loss = (adversarial_loss(validity1, valid) + adversarial_loss(validity2, valid)) / 2\n\n g_loss.backward()\n optimizer_G.step()\n\n # ----------------------\n # Train Discriminators\n # ----------------------\n\n optimizer_D.zero_grad()\n\n # Determine validity of real and generated images\n validity1_real, validity2_real = coupled_discriminators(imgs1, imgs2)\n validity1_fake, validity2_fake = coupled_discriminators(gen_imgs1.detach(), gen_imgs2.detach())\n\n d_loss = (\n adversarial_loss(validity1_real, valid)\n + adversarial_loss(validity1_fake, fake)\n + adversarial_loss(validity2_real, valid)\n + adversarial_loss(validity2_fake, fake)\n ) / 4\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader1), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader1) + i\n if batches_done % opt.sample_interval == 0:\n gen_imgs = torch.cat((gen_imgs1.data, gen_imgs2.data), 0)\n save_image(gen_imgs, \"images/%d.png\" % batches_done, nrow=8, normalize=True)\n"
},
{
"path": "implementations/cogan/mnistm.py",
"content": "\"\"\"Dataset setting and data loader for MNIST-M.\n\nModified from\nhttps://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py\n\nCREDIT: https://github.com/corenel\n\"\"\"\n\nfrom __future__ import print_function\n\nimport errno\nimport os\n\nimport torch\nimport torch.utils.data as data\nfrom PIL import Image\n\n\nclass MNISTM(data.Dataset):\n \"\"\"`MNIST-M Dataset.\"\"\"\n\n url = \"https://github.com/VanushVaswani/keras_mnistm/releases/download/1.0/keras_mnistm.pkl.gz\"\n\n raw_folder = \"raw\"\n processed_folder = \"processed\"\n training_file = \"mnist_m_train.pt\"\n test_file = \"mnist_m_test.pt\"\n\n def __init__(self, root, mnist_root=\"data\", train=True, transform=None, target_transform=None, download=False):\n \"\"\"Init MNIST-M dataset.\"\"\"\n super(MNISTM, self).__init__()\n self.root = os.path.expanduser(root)\n self.mnist_root = os.path.expanduser(mnist_root)\n self.transform = transform\n self.target_transform = target_transform\n self.train = train # training set or test set\n\n if download:\n self.download()\n\n if not self._check_exists():\n raise RuntimeError(\"Dataset not found.\" + \" You can use download=True to download it\")\n\n if self.train:\n self.train_data, self.train_labels = torch.load(\n os.path.join(self.root, self.processed_folder, self.training_file)\n )\n else:\n self.test_data, self.test_labels = torch.load(\n os.path.join(self.root, self.processed_folder, self.test_file)\n )\n\n def __getitem__(self, index):\n \"\"\"Get images and target for data loader.\n\n Args:\n index (int): Index\n\n Returns:\n tuple: (image, target) where target is index of the target class.\n \"\"\"\n if self.train:\n img, target = self.train_data[index], self.train_labels[index]\n else:\n img, target = self.test_data[index], self.test_labels[index]\n\n # doing this so that it is consistent with all other datasets\n # to return a PIL Image\n img = Image.fromarray(img.squeeze().numpy(), mode=\"RGB\")\n\n if self.transform is not None:\n img = self.transform(img)\n\n if self.target_transform is not None:\n target = self.target_transform(target)\n\n return img, target\n\n def __len__(self):\n \"\"\"Return size of dataset.\"\"\"\n if self.train:\n return len(self.train_data)\n else:\n return len(self.test_data)\n\n def _check_exists(self):\n return os.path.exists(os.path.join(self.root, self.processed_folder, self.training_file)) and os.path.exists(\n os.path.join(self.root, self.processed_folder, self.test_file)\n )\n\n def download(self):\n \"\"\"Download the MNIST data.\"\"\"\n # import essential packages\n from six.moves import urllib\n import gzip\n import pickle\n from torchvision import datasets\n\n # check if dataset already exists\n if self._check_exists():\n return\n\n # make data dirs\n try:\n os.makedirs(os.path.join(self.root, self.raw_folder))\n os.makedirs(os.path.join(self.root, self.processed_folder))\n except OSError as e:\n if e.errno == errno.EEXIST:\n pass\n else:\n raise\n\n # download pkl files\n print(\"Downloading \" + self.url)\n filename = self.url.rpartition(\"/\")[2]\n file_path = os.path.join(self.root, self.raw_folder, filename)\n if not os.path.exists(file_path.replace(\".gz\", \"\")):\n data = urllib.request.urlopen(self.url)\n with open(file_path, \"wb\") as f:\n f.write(data.read())\n with open(file_path.replace(\".gz\", \"\"), \"wb\") as out_f, gzip.GzipFile(file_path) as zip_f:\n out_f.write(zip_f.read())\n os.unlink(file_path)\n\n # process and save as torch files\n print(\"Processing...\")\n\n # load MNIST-M images from pkl file\n with open(file_path.replace(\".gz\", \"\"), \"rb\") as f:\n mnist_m_data = pickle.load(f, encoding=\"bytes\")\n mnist_m_train_data = torch.ByteTensor(mnist_m_data[b\"train\"])\n mnist_m_test_data = torch.ByteTensor(mnist_m_data[b\"test\"])\n\n # get MNIST labels\n mnist_train_labels = datasets.MNIST(root=self.mnist_root, train=True, download=True).train_labels\n mnist_test_labels = datasets.MNIST(root=self.mnist_root, train=False, download=True).test_labels\n\n # save MNIST-M dataset\n training_set = (mnist_m_train_data, mnist_train_labels)\n test_set = (mnist_m_test_data, mnist_test_labels)\n with open(os.path.join(self.root, self.processed_folder, self.training_file), \"wb\") as f:\n torch.save(training_set, f)\n with open(os.path.join(self.root, self.processed_folder, self.test_file), \"wb\") as f:\n torch.save(test_set, f)\n\n print(\"Done!\")\n"
},
{
"path": "implementations/context_encoder/context_encoder.py",
"content": "\"\"\"\nInpainting using Generative Adversarial Networks.\nThe dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0\n(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)\nInstrustion on running the script:\n1. Download the dataset from the provided link\n2. Save the folder 'img_align_celeba' to '../../data/'\n4. Run the sript using command 'python3 context_encoder.py'\n\"\"\"\n\nimport argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\nfrom PIL import Image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom datasets import *\nfrom models import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=8, help=\"size of the batches\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"img_align_celeba\", help=\"name of the dataset\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=4, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=128, help=\"size of each image dimension\")\nparser.add_argument(\"--mask_size\", type=int, default=64, help=\"size of random mask\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=500, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n# Calculate output of image discriminator (PatchGAN)\npatch_h, patch_w = int(opt.mask_size / 2 ** 3), int(opt.mask_size / 2 ** 3)\npatch = (1, patch_h, patch_w)\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n# Loss function\nadversarial_loss = torch.nn.MSELoss()\npixelwise_loss = torch.nn.L1Loss()\n\n# Initialize generator and discriminator\ngenerator = Generator(channels=opt.channels)\ndiscriminator = Discriminator(channels=opt.channels)\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n pixelwise_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Dataset loader\ntransforms_ = [\n transforms.Resize((opt.img_size, opt.img_size), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\ntest_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, mode=\"val\"),\n batch_size=12,\n shuffle=True,\n num_workers=1,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef save_sample(batches_done):\n samples, masked_samples, i = next(iter(test_dataloader))\n samples = Variable(samples.type(Tensor))\n masked_samples = Variable(masked_samples.type(Tensor))\n i = i[0].item() # Upper-left coordinate of mask\n # Generate inpainted image\n gen_mask = generator(masked_samples)\n filled_samples = masked_samples.clone()\n filled_samples[:, :, i : i + opt.mask_size, i : i + opt.mask_size] = gen_mask\n # Save sample\n sample = torch.cat((masked_samples.data, filled_samples.data, samples.data), -2)\n save_image(sample, \"images/%d.png\" % batches_done, nrow=6, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, masked_imgs, masked_parts) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], *patch).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], *patch).fill_(0.0), requires_grad=False)\n\n # Configure input\n imgs = Variable(imgs.type(Tensor))\n masked_imgs = Variable(masked_imgs.type(Tensor))\n masked_parts = Variable(masked_parts.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Generate a batch of images\n gen_parts = generator(masked_imgs)\n\n # Adversarial and pixelwise loss\n g_adv = adversarial_loss(discriminator(gen_parts), valid)\n g_pixel = pixelwise_loss(gen_parts, masked_parts)\n # Total loss\n g_loss = 0.001 * g_adv + 0.999 * g_pixel\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(masked_parts), valid)\n fake_loss = adversarial_loss(discriminator(gen_parts.detach()), fake)\n d_loss = 0.5 * (real_loss + fake_loss)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G adv: %f, pixel: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_adv.item(), g_pixel.item())\n )\n\n # Generate sample at sample interval\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_sample(batches_done)\n"
},
{
"path": "implementations/context_encoder/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, img_size=128, mask_size=64, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n self.img_size = img_size\n self.mask_size = mask_size\n self.mode = mode\n self.files = sorted(glob.glob(\"%s/*.jpg\" % root))\n self.files = self.files[:-4000] if mode == \"train\" else self.files[-4000:]\n\n def apply_random_mask(self, img):\n \"\"\"Randomly masks image\"\"\"\n y1, x1 = np.random.randint(0, self.img_size - self.mask_size, 2)\n y2, x2 = y1 + self.mask_size, x1 + self.mask_size\n masked_part = img[:, y1:y2, x1:x2]\n masked_img = img.clone()\n masked_img[:, y1:y2, x1:x2] = 1\n\n return masked_img, masked_part\n\n def apply_center_mask(self, img):\n \"\"\"Mask center part of image\"\"\"\n # Get upper-left pixel coordinate\n i = (self.img_size - self.mask_size) // 2\n masked_img = img.clone()\n masked_img[:, i : i + self.mask_size, i : i + self.mask_size] = 1\n\n return masked_img, i\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n img = self.transform(img)\n if self.mode == \"train\":\n # For training data perform random mask\n masked_img, aux = self.apply_random_mask(img)\n else:\n # For test data mask the center of the image\n masked_img, aux = self.apply_center_mask(img)\n\n return img, masked_img, aux\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/context_encoder/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\nclass Generator(nn.Module):\n def __init__(self, channels=3):\n super(Generator, self).__init__()\n\n def downsample(in_feat, out_feat, normalize=True):\n layers = [nn.Conv2d(in_feat, out_feat, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.BatchNorm2d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2))\n return layers\n\n def upsample(in_feat, out_feat, normalize=True):\n layers = [nn.ConvTranspose2d(in_feat, out_feat, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.BatchNorm2d(out_feat, 0.8))\n layers.append(nn.ReLU())\n return layers\n\n self.model = nn.Sequential(\n *downsample(channels, 64, normalize=False),\n *downsample(64, 64),\n *downsample(64, 128),\n *downsample(128, 256),\n *downsample(256, 512),\n nn.Conv2d(512, 4000, 1),\n *upsample(4000, 512),\n *upsample(512, 256),\n *upsample(256, 128),\n *upsample(128, 64),\n nn.Conv2d(64, channels, 3, 1, 1),\n nn.Tanh()\n )\n\n def forward(self, x):\n return self.model(x)\n\n\nclass Discriminator(nn.Module):\n def __init__(self, channels=3):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, stride, normalize):\n \"\"\"Returns layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 3, stride, 1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n layers = []\n in_filters = channels\n for out_filters, stride, normalize in [(64, 2, False), (128, 2, True), (256, 2, True), (512, 1, True)]:\n layers.extend(discriminator_block(in_filters, out_filters, stride, normalize))\n in_filters = out_filters\n\n layers.append(nn.Conv2d(out_filters, 1, 3, 1, 1))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/cyclegan/cyclegan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport datetime\nimport time\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image, make_grid\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\nfrom utils import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"monet2photo\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=1, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=256, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=256, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving generator outputs\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between saving model checkpoints\")\nparser.add_argument(\"--n_residual_blocks\", type=int, default=9, help=\"number of residual blocks in generator\")\nparser.add_argument(\"--lambda_cyc\", type=float, default=10.0, help=\"cycle loss weight\")\nparser.add_argument(\"--lambda_id\", type=float, default=5.0, help=\"identity loss weight\")\nopt = parser.parse_args()\nprint(opt)\n\n# Create sample and checkpoint directories\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\n# Losses\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_cycle = torch.nn.L1Loss()\ncriterion_identity = torch.nn.L1Loss()\n\ncuda = torch.cuda.is_available()\n\ninput_shape = (opt.channels, opt.img_height, opt.img_width)\n\n# Initialize generator and discriminator\nG_AB = GeneratorResNet(input_shape, opt.n_residual_blocks)\nG_BA = GeneratorResNet(input_shape, opt.n_residual_blocks)\nD_A = Discriminator(input_shape)\nD_B = Discriminator(input_shape)\n\nif cuda:\n G_AB = G_AB.cuda()\n G_BA = G_BA.cuda()\n D_A = D_A.cuda()\n D_B = D_B.cuda()\n criterion_GAN.cuda()\n criterion_cycle.cuda()\n criterion_identity.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n G_AB.load_state_dict(torch.load(\"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, opt.epoch)))\n G_BA.load_state_dict(torch.load(\"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_A.load_state_dict(torch.load(\"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_B.load_state_dict(torch.load(\"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n G_AB.apply(weights_init_normal)\n G_BA.apply(weights_init_normal)\n D_A.apply(weights_init_normal)\n D_B.apply(weights_init_normal)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\noptimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Learning rate update schedulers\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n optimizer_G, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D_A, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D_B, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\n# Buffers of previously generated samples\nfake_A_buffer = ReplayBuffer()\nfake_B_buffer = ReplayBuffer()\n\n# Image transformations\ntransforms_ = [\n transforms.Resize(int(opt.img_height * 1.12), Image.BICUBIC),\n transforms.RandomCrop((opt.img_height, opt.img_width)),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\n# Training data loader\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n# Test data loader\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True, mode=\"test\"),\n batch_size=5,\n shuffle=True,\n num_workers=1,\n)\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the test set\"\"\"\n imgs = next(iter(val_dataloader))\n G_AB.eval()\n G_BA.eval()\n real_A = Variable(imgs[\"A\"].type(Tensor))\n fake_B = G_AB(real_A)\n real_B = Variable(imgs[\"B\"].type(Tensor))\n fake_A = G_BA(real_B)\n # Arange images along x-axis\n real_A = make_grid(real_A, nrow=5, normalize=True)\n real_B = make_grid(real_B, nrow=5, normalize=True)\n fake_A = make_grid(fake_A, nrow=5, normalize=True)\n fake_B = make_grid(fake_B, nrow=5, normalize=True)\n # Arange images along y-axis\n image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)\n save_image(image_grid, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), normalize=False)\n\n\n# ----------\n# Training\n# ----------\n\nprev_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Set model input\n real_A = Variable(batch[\"A\"].type(Tensor))\n real_B = Variable(batch[\"B\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((real_A.size(0), *D_A.output_shape))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((real_A.size(0), *D_A.output_shape))), requires_grad=False)\n\n # ------------------\n # Train Generators\n # ------------------\n\n G_AB.train()\n G_BA.train()\n\n optimizer_G.zero_grad()\n\n # Identity loss\n loss_id_A = criterion_identity(G_BA(real_A), real_A)\n loss_id_B = criterion_identity(G_AB(real_B), real_B)\n\n loss_identity = (loss_id_A + loss_id_B) / 2\n\n # GAN loss\n fake_B = G_AB(real_A)\n loss_GAN_AB = criterion_GAN(D_B(fake_B), valid)\n fake_A = G_BA(real_B)\n loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)\n\n loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2\n\n # Cycle loss\n recov_A = G_BA(fake_B)\n loss_cycle_A = criterion_cycle(recov_A, real_A)\n recov_B = G_AB(fake_A)\n loss_cycle_B = criterion_cycle(recov_B, real_B)\n\n loss_cycle = (loss_cycle_A + loss_cycle_B) / 2\n\n # Total loss\n loss_G = loss_GAN + opt.lambda_cyc * loss_cycle + opt.lambda_id * loss_identity\n\n loss_G.backward()\n optimizer_G.step()\n\n # -----------------------\n # Train Discriminator A\n # -----------------------\n\n optimizer_D_A.zero_grad()\n\n # Real loss\n loss_real = criterion_GAN(D_A(real_A), valid)\n # Fake loss (on batch of previously generated samples)\n fake_A_ = fake_A_buffer.push_and_pop(fake_A)\n loss_fake = criterion_GAN(D_A(fake_A_.detach()), fake)\n # Total loss\n loss_D_A = (loss_real + loss_fake) / 2\n\n loss_D_A.backward()\n optimizer_D_A.step()\n\n # -----------------------\n # Train Discriminator B\n # -----------------------\n\n optimizer_D_B.zero_grad()\n\n # Real loss\n loss_real = criterion_GAN(D_B(real_B), valid)\n # Fake loss (on batch of previously generated samples)\n fake_B_ = fake_B_buffer.push_and_pop(fake_B)\n loss_fake = criterion_GAN(D_B(fake_B_.detach()), fake)\n # Total loss\n loss_D_B = (loss_real + loss_fake) / 2\n\n loss_D_B.backward()\n optimizer_D_B.step()\n\n loss_D = (loss_D_A + loss_D_B) / 2\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, adv: %f, cycle: %f, identity: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D.item(),\n loss_G.item(),\n loss_GAN.item(),\n loss_cycle.item(),\n loss_identity.item(),\n time_left,\n )\n )\n\n # If at sample interval save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n # Update learning rates\n lr_scheduler_G.step()\n lr_scheduler_D_A.step()\n lr_scheduler_D_B.step()\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(G_AB.state_dict(), \"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(G_BA.state_dict(), \"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_A.state_dict(), \"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_B.state_dict(), \"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/cyclegan/datasets.py",
"content": "import glob\nimport random\nimport os\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\ndef to_rgb(image):\n rgb_image = Image.new(\"RGB\", image.size)\n rgb_image.paste(image)\n return rgb_image\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, unaligned=False, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n self.unaligned = unaligned\n\n self.files_A = sorted(glob.glob(os.path.join(root, \"%s/A\" % mode) + \"/*.*\"))\n self.files_B = sorted(glob.glob(os.path.join(root, \"%s/B\" % mode) + \"/*.*\"))\n\n def __getitem__(self, index):\n image_A = Image.open(self.files_A[index % len(self.files_A)])\n\n if self.unaligned:\n image_B = Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)])\n else:\n image_B = Image.open(self.files_B[index % len(self.files_B)])\n\n # Convert grayscale images to rgb\n if image_A.mode != \"RGB\":\n image_A = to_rgb(image_A)\n if image_B.mode != \"RGB\":\n image_B = to_rgb(image_B)\n\n item_A = self.transform(image_A)\n item_B = self.transform(image_B)\n return {\"A\": item_A, \"B\": item_B}\n\n def __len__(self):\n return max(len(self.files_A), len(self.files_B))\n"
},
{
"path": "implementations/cyclegan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n if hasattr(m, \"bias\") and m.bias is not None:\n torch.nn.init.constant_(m.bias.data, 0.0)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n# RESNET\n##############################\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, in_features):\n super(ResidualBlock, self).__init__()\n\n self.block = nn.Sequential(\n nn.ReflectionPad2d(1),\n nn.Conv2d(in_features, in_features, 3),\n nn.InstanceNorm2d(in_features),\n nn.ReLU(inplace=True),\n nn.ReflectionPad2d(1),\n nn.Conv2d(in_features, in_features, 3),\n nn.InstanceNorm2d(in_features),\n )\n\n def forward(self, x):\n return x + self.block(x)\n\n\nclass GeneratorResNet(nn.Module):\n def __init__(self, input_shape, num_residual_blocks):\n super(GeneratorResNet, self).__init__()\n\n channels = input_shape[0]\n\n # Initial convolution block\n out_features = 64\n model = [\n nn.ReflectionPad2d(channels),\n nn.Conv2d(channels, out_features, 7),\n nn.InstanceNorm2d(out_features),\n nn.ReLU(inplace=True),\n ]\n in_features = out_features\n\n # Downsampling\n for _ in range(2):\n out_features *= 2\n model += [\n nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n nn.InstanceNorm2d(out_features),\n nn.ReLU(inplace=True),\n ]\n in_features = out_features\n\n # Residual blocks\n for _ in range(num_residual_blocks):\n model += [ResidualBlock(out_features)]\n\n # Upsampling\n for _ in range(2):\n out_features //= 2\n model += [\n nn.Upsample(scale_factor=2),\n nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),\n nn.InstanceNorm2d(out_features),\n nn.ReLU(inplace=True),\n ]\n in_features = out_features\n\n # Output layer\n model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x):\n return self.model(x)\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n\n channels, height, width = input_shape\n\n # Calculate output shape of image discriminator (PatchGAN)\n self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)\n\n def discriminator_block(in_filters, out_filters, normalize=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *discriminator_block(channels, 64, normalize=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n *discriminator_block(256, 512),\n nn.ZeroPad2d((1, 0, 1, 0)),\n nn.Conv2d(512, 1, 4, padding=1)\n )\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/cyclegan/utils.py",
"content": "import random\nimport time\nimport datetime\nimport sys\n\nfrom torch.autograd import Variable\nimport torch\nimport numpy as np\n\nfrom torchvision.utils import save_image\n\n\nclass ReplayBuffer:\n def __init__(self, max_size=50):\n assert max_size > 0, \"Empty buffer or trying to create a black hole. Be careful.\"\n self.max_size = max_size\n self.data = []\n\n def push_and_pop(self, data):\n to_return = []\n for element in data.data:\n element = torch.unsqueeze(element, 0)\n if len(self.data) < self.max_size:\n self.data.append(element)\n to_return.append(element)\n else:\n if random.uniform(0, 1) > 0.5:\n i = random.randint(0, self.max_size - 1)\n to_return.append(self.data[i].clone())\n self.data[i] = element\n else:\n to_return.append(element)\n return Variable(torch.cat(to_return))\n\n\nclass LambdaLR:\n def __init__(self, n_epochs, offset, decay_start_epoch):\n assert (n_epochs - decay_start_epoch) > 0, \"Decay must start before the training session ends!\"\n self.n_epochs = n_epochs\n self.offset = offset\n self.decay_start_epoch = decay_start_epoch\n\n def step(self, epoch):\n return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)\n"
},
{
"path": "implementations/dcgan/dcgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, z):\n out = self.l1(z)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.model = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())\n\n def forward(self, img):\n out = self.model(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n\n return validity\n\n\n# Loss function\nadversarial_loss = torch.nn.BCELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(real_imgs), valid)\n fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/discogan/datasets.py",
"content": "import glob\nimport os\nimport torch\n\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, mode='train'):\n self.transform = transforms.Compose(transforms_)\n\n self.files = sorted(glob.glob(os.path.join(root, mode) + '/*.*'))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n w, h = img.size\n img_A = img.crop((0, 0, w/2, h))\n img_B = img.crop((w/2, 0, w, h))\n\n if np.random.random() < 0.5:\n img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], 'RGB')\n img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], 'RGB')\n\n img_A = self.transform(img_A)\n img_B = self.transform(img_B)\n\n return {'A': img_A, 'B': img_B}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/discogan/discogan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport sys\nimport datetime\nimport time\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"edges2shoes\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=64, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=64, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving generator samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nopt = parser.parse_args()\nprint(opt)\n\n# Create sample and checkpoint directories\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\n# Losses\nadversarial_loss = torch.nn.MSELoss()\ncycle_loss = torch.nn.L1Loss()\npixelwise_loss = torch.nn.L1Loss()\n\ncuda = torch.cuda.is_available()\n\ninput_shape = (opt.channels, opt.img_height, opt.img_width)\n\n# Initialize generator and discriminator\nG_AB = GeneratorUNet(input_shape)\nG_BA = GeneratorUNet(input_shape)\nD_A = Discriminator(input_shape)\nD_B = Discriminator(input_shape)\n\nif cuda:\n G_AB = G_AB.cuda()\n G_BA = G_BA.cuda()\n D_A = D_A.cuda()\n D_B = D_B.cuda()\n adversarial_loss.cuda()\n cycle_loss.cuda()\n pixelwise_loss.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n G_AB.load_state_dict(torch.load(\"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, opt.epoch)))\n G_BA.load_state_dict(torch.load(\"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_A.load_state_dict(torch.load(\"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_B.load_state_dict(torch.load(\"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n G_AB.apply(weights_init_normal)\n G_BA.apply(weights_init_normal)\n D_A.apply(weights_init_normal)\n D_B.apply(weights_init_normal)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\noptimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Input tensor type\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\n# Dataset loader\ntransforms_ = [\n transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, mode=\"train\"),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, mode=\"val\"),\n batch_size=16,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the validation set\"\"\"\n imgs = next(iter(val_dataloader))\n G_AB.eval()\n G_BA.eval()\n real_A = Variable(imgs[\"A\"].type(Tensor))\n fake_B = G_AB(real_A)\n real_B = Variable(imgs[\"B\"].type(Tensor))\n fake_A = G_BA(real_B)\n img_sample = torch.cat((real_A.data, fake_B.data, real_B.data, fake_A.data), 0)\n save_image(img_sample, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=8, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nprev_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Model inputs\n real_A = Variable(batch[\"A\"].type(Tensor))\n real_B = Variable(batch[\"B\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((real_A.size(0), *D_A.output_shape))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((real_A.size(0), *D_A.output_shape))), requires_grad=False)\n\n # ------------------\n # Train Generators\n # ------------------\n\n G_AB.train()\n G_BA.train()\n\n optimizer_G.zero_grad()\n\n # GAN loss\n fake_B = G_AB(real_A)\n loss_GAN_AB = adversarial_loss(D_B(fake_B), valid)\n fake_A = G_BA(real_B)\n loss_GAN_BA = adversarial_loss(D_A(fake_A), valid)\n\n loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2\n\n # Pixelwise translation loss\n loss_pixelwise = (pixelwise_loss(fake_A, real_A) + pixelwise_loss(fake_B, real_B)) / 2\n\n # Cycle loss\n loss_cycle_A = cycle_loss(G_BA(fake_B), real_A)\n loss_cycle_B = cycle_loss(G_AB(fake_A), real_B)\n loss_cycle = (loss_cycle_A + loss_cycle_B) / 2\n\n # Total loss\n loss_G = loss_GAN + loss_cycle + loss_pixelwise\n\n loss_G.backward()\n optimizer_G.step()\n\n # -----------------------\n # Train Discriminator A\n # -----------------------\n\n optimizer_D_A.zero_grad()\n\n # Real loss\n loss_real = adversarial_loss(D_A(real_A), valid)\n # Fake loss (on batch of previously generated samples)\n loss_fake = adversarial_loss(D_A(fake_A.detach()), fake)\n # Total loss\n loss_D_A = (loss_real + loss_fake) / 2\n\n loss_D_A.backward()\n optimizer_D_A.step()\n\n # -----------------------\n # Train Discriminator B\n # -----------------------\n\n optimizer_D_B.zero_grad()\n # Real loss\n loss_real = adversarial_loss(D_B(real_B), valid)\n # Fake loss (on batch of previously generated samples)\n loss_fake = adversarial_loss(D_B(fake_B.detach()), fake)\n # Total loss\n loss_D_B = (loss_real + loss_fake) / 2\n\n loss_D_B.backward()\n optimizer_D_B.step()\n\n loss_D = 0.5 * (loss_D_A + loss_D_B)\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, adv: %f, pixel: %f, cycle: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D.item(),\n loss_G.item(),\n loss_GAN.item(),\n loss_pixelwise.item(),\n loss_cycle.item(),\n time_left,\n )\n )\n\n # If at sample interval save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(G_AB.state_dict(), \"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(G_BA.state_dict(), \"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_A.state_dict(), \"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_B.state_dict(), \"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/discogan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n# U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n super(UNetDown, self).__init__()\n layers = [nn.Conv2d(in_size, out_size, 4, 2, 1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_size))\n layers.append(nn.LeakyReLU(0.2))\n if dropout:\n layers.append(nn.Dropout(dropout))\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass UNetUp(nn.Module):\n def __init__(self, in_size, out_size, dropout=0.0):\n super(UNetUp, self).__init__()\n layers = [nn.ConvTranspose2d(in_size, out_size, 4, 2, 1), nn.InstanceNorm2d(out_size), nn.ReLU(inplace=True)]\n if dropout:\n layers.append(nn.Dropout(dropout))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x, skip_input):\n x = self.model(x)\n x = torch.cat((x, skip_input), 1)\n\n return x\n\n\nclass GeneratorUNet(nn.Module):\n def __init__(self, input_shape):\n super(GeneratorUNet, self).__init__()\n channels, _, _ = input_shape\n self.down1 = UNetDown(channels, 64, normalize=False)\n self.down2 = UNetDown(64, 128)\n self.down3 = UNetDown(128, 256, dropout=0.5)\n self.down4 = UNetDown(256, 512, dropout=0.5)\n self.down5 = UNetDown(512, 512, dropout=0.5)\n self.down6 = UNetDown(512, 512, dropout=0.5, normalize=False)\n\n self.up1 = UNetUp(512, 512, dropout=0.5)\n self.up2 = UNetUp(1024, 512, dropout=0.5)\n self.up3 = UNetUp(1024, 256, dropout=0.5)\n self.up4 = UNetUp(512, 128)\n self.up5 = UNetUp(256, 64)\n\n self.final = nn.Sequential(\n nn.Upsample(scale_factor=2), nn.ZeroPad2d((1, 0, 1, 0)), nn.Conv2d(128, channels, 4, padding=1), nn.Tanh()\n )\n\n def forward(self, x):\n # U-Net generator with skip connections from encoder to decoder\n d1 = self.down1(x)\n d2 = self.down2(d1)\n d3 = self.down3(d2)\n d4 = self.down4(d3)\n d5 = self.down5(d4)\n d6 = self.down6(d5)\n u1 = self.up1(d6, d5)\n u2 = self.up2(u1, d4)\n u3 = self.up3(u2, d3)\n u4 = self.up4(u3, d2)\n u5 = self.up5(u4, d1)\n\n return self.final(u5)\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n\n channels, height, width = input_shape\n # Calculate output of image discriminator (PatchGAN)\n self.output_shape = (1, height // 2 ** 3, width // 2 ** 3)\n\n def discriminator_block(in_filters, out_filters, normalization=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalization:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *discriminator_block(channels, 64, normalization=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n nn.ZeroPad2d((1, 0, 1, 0)),\n nn.Conv2d(256, 1, 4, padding=1)\n )\n\n def forward(self, img):\n # Concatenate image and condition image by channels to produce input\n return self.model(img)\n"
},
{
"path": "implementations/dragan/dragan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\nimport torch.autograd as autograd\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=1000, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise):\n out = self.l1(noise)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.model = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())\n\n def forward(self, img):\n out = self.model(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n\n return validity\n\n\n# Loss function\nadversarial_loss = torch.nn.BCELoss()\n\n# Loss weight for gradient penalty\nlambda_gp = 10\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef compute_gradient_penalty(D, X):\n \"\"\"Calculates the gradient penalty loss for DRAGAN\"\"\"\n # Random weight term for interpolation\n alpha = Tensor(np.random.random(size=X.shape))\n\n interpolates = alpha * X + ((1 - alpha) * (X + 0.5 * X.std() * torch.rand(X.size())))\n interpolates = Variable(interpolates, requires_grad=True)\n\n d_interpolates = D(interpolates)\n\n fake = Variable(Tensor(X.shape[0], 1).fill_(1.0), requires_grad=False)\n\n # Get gradient w.r.t. interpolates\n gradients = autograd.grad(\n outputs=d_interpolates,\n inputs=interpolates,\n grad_outputs=fake,\n create_graph=True,\n retain_graph=True,\n only_inputs=True,\n )[0]\n\n gradient_penalty = lambda_gp * ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n return gradient_penalty\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(mnist_loader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(real_imgs), valid)\n fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = (real_loss + fake_loss) / 2\n\n # Calculate gradient penalty\n gradient_penalty = compute_gradient_penalty(discriminator, real_imgs.data)\n gradient_penalty.backward()\n\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(mnist_loader), d_loss.item(), g_loss.item())\n )\n\n save_image(gen_imgs.data, \"images/%d.png\" % epoch, nrow=int(math.sqrt(opt.batch_size)), normalize=True)\n"
},
{
"path": "implementations/dualgan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n\n self.files = sorted(glob.glob(os.path.join(root, mode) + \"/*.*\"))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n w, h = img.size\n img_A = img.crop((0, 0, w / 2, h))\n img_B = img.crop((w / 2, 0, w, h))\n\n if np.random.random() < 0.5:\n img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\")\n img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], \"RGB\")\n\n img_A = self.transform(img_A)\n img_B = self.transform(img_B)\n\n return {\"A\": img_A, \"B\": img_B}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/dualgan/dualgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport scipy\nimport sys\nimport time\nimport datetime\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\nimport torch.autograd as autograd\n\nfrom datasets import *\nfrom models import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=8, help=\"size of the batches\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"edges2shoes\", help=\"name of the dataset\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_size\", type=int, default=128, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--n_critic\", type=int, default=5, help=\"number of training steps for discriminator per iter\")\nparser.add_argument(\"--sample_interval\", type=int, default=200, help=\"interval betwen image samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nopt = parser.parse_args()\nprint(opt)\n\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n# Loss function\ncycle_loss = torch.nn.L1Loss()\n\n# Loss weights\nlambda_adv = 1\nlambda_cycle = 10\nlambda_gp = 10\n\n# Initialize generator and discriminator\nG_AB = Generator()\nG_BA = Generator()\nD_A = Discriminator()\nD_B = Discriminator()\n\nif cuda:\n G_AB.cuda()\n G_BA.cuda()\n D_A.cuda()\n D_B.cuda()\n cycle_loss.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n G_AB.load_state_dict(torch.load(\"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, opt.epoch)))\n G_BA.load_state_dict(torch.load(\"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_A.load_state_dict(torch.load(\"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D_B.load_state_dict(torch.load(\"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n G_AB.apply(weights_init_normal)\n G_BA.apply(weights_init_normal)\n D_A.apply(weights_init_normal)\n D_B.apply(weights_init_normal)\n\n# Configure data loader\ntransforms_ = [\n transforms.Resize((opt.img_size, opt.img_size), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, mode=\"val\", transforms_=transforms_),\n batch_size=16,\n shuffle=True,\n num_workers=1,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\noptimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n\ndef compute_gradient_penalty(D, real_samples, fake_samples):\n \"\"\"Calculates the gradient penalty loss for WGAN GP\"\"\"\n # Random weight term for interpolation between real and fake samples\n alpha = FloatTensor(np.random.random((real_samples.size(0), 1, 1, 1)))\n # Get random interpolation between real and fake samples\n interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)\n validity = D(interpolates)\n fake = Variable(FloatTensor(np.ones(validity.shape)), requires_grad=False)\n # Get gradient w.r.t. interpolates\n gradients = autograd.grad(\n outputs=validity,\n inputs=interpolates,\n grad_outputs=fake,\n create_graph=True,\n retain_graph=True,\n only_inputs=True,\n )[0]\n gradients = gradients.view(gradients.size(0), -1)\n gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n return gradient_penalty\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the test set\"\"\"\n imgs = next(iter(val_dataloader))\n real_A = Variable(imgs[\"A\"].type(FloatTensor))\n fake_B = G_AB(real_A)\n AB = torch.cat((real_A.data, fake_B.data), -2)\n real_B = Variable(imgs[\"B\"].type(FloatTensor))\n fake_A = G_BA(real_B)\n BA = torch.cat((real_B.data, fake_A.data), -2)\n img_sample = torch.cat((AB, BA), 0)\n save_image(img_sample, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=8, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nbatches_done = 0\nprev_time = time.time()\nfor epoch in range(opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Configure input\n imgs_A = Variable(batch[\"A\"].type(FloatTensor))\n imgs_B = Variable(batch[\"B\"].type(FloatTensor))\n\n # ----------------------\n # Train Discriminators\n # ----------------------\n\n optimizer_D_A.zero_grad()\n optimizer_D_B.zero_grad()\n\n # Generate a batch of images\n fake_A = G_BA(imgs_B).detach()\n fake_B = G_AB(imgs_A).detach()\n\n # ----------\n # Domain A\n # ----------\n\n # Compute gradient penalty for improved wasserstein training\n gp_A = compute_gradient_penalty(D_A, imgs_A.data, fake_A.data)\n # Adversarial loss\n D_A_loss = -torch.mean(D_A(imgs_A)) + torch.mean(D_A(fake_A)) + lambda_gp * gp_A\n\n # ----------\n # Domain B\n # ----------\n\n # Compute gradient penalty for improved wasserstein training\n gp_B = compute_gradient_penalty(D_B, imgs_B.data, fake_B.data)\n # Adversarial loss\n D_B_loss = -torch.mean(D_B(imgs_B)) + torch.mean(D_B(fake_B)) + lambda_gp * gp_B\n\n # Total loss\n D_loss = D_A_loss + D_B_loss\n\n D_loss.backward()\n optimizer_D_A.step()\n optimizer_D_B.step()\n\n if i % opt.n_critic == 0:\n\n # ------------------\n # Train Generators\n # ------------------\n\n optimizer_G.zero_grad()\n\n # Translate images to opposite domain\n fake_A = G_BA(imgs_B)\n fake_B = G_AB(imgs_A)\n\n # Reconstruct images\n recov_A = G_BA(fake_B)\n recov_B = G_AB(fake_A)\n\n # Adversarial loss\n G_adv = -torch.mean(D_A(fake_A)) - torch.mean(D_B(fake_B))\n # Cycle loss\n G_cycle = cycle_loss(recov_A, imgs_A) + cycle_loss(recov_B, imgs_B)\n # Total loss\n G_loss = lambda_adv * G_adv + lambda_cycle * G_cycle\n\n G_loss.backward()\n optimizer_G.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time) / opt.n_critic)\n prev_time = time.time()\n\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, cycle: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n D_loss.item(),\n G_adv.data.item(),\n G_cycle.item(),\n time_left,\n )\n )\n\n # Check sample interval => save sample if there\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n batches_done += 1\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(G_AB.state_dict(), \"saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(G_BA.state_dict(), \"saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_A.state_dict(), \"saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D_B.state_dict(), \"saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/dualgan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torchvision.models import vgg19\nimport math\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n# U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n super(UNetDown, self).__init__()\n layers = [nn.Conv2d(in_size, out_size, 4, stride=2, padding=1, bias=False)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_size, affine=True))\n layers.append(nn.LeakyReLU(0.2))\n if dropout:\n layers.append(nn.Dropout(dropout))\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass UNetUp(nn.Module):\n def __init__(self, in_size, out_size, dropout=0.0):\n super(UNetUp, self).__init__()\n layers = [\n nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),\n nn.InstanceNorm2d(out_size, affine=True),\n nn.ReLU(inplace=True),\n ]\n if dropout:\n layers.append(nn.Dropout(dropout))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x, skip_input):\n x = self.model(x)\n x = torch.cat((x, skip_input), 1)\n\n return x\n\n\nclass Generator(nn.Module):\n def __init__(self, channels=3):\n super(Generator, self).__init__()\n\n self.down1 = UNetDown(channels, 64, normalize=False)\n self.down2 = UNetDown(64, 128)\n self.down3 = UNetDown(128, 256)\n self.down4 = UNetDown(256, 512, dropout=0.5)\n self.down5 = UNetDown(512, 512, dropout=0.5)\n self.down6 = UNetDown(512, 512, dropout=0.5)\n self.down7 = UNetDown(512, 512, dropout=0.5, normalize=False)\n\n self.up1 = UNetUp(512, 512, dropout=0.5)\n self.up2 = UNetUp(1024, 512, dropout=0.5)\n self.up3 = UNetUp(1024, 512, dropout=0.5)\n self.up4 = UNetUp(1024, 256)\n self.up5 = UNetUp(512, 128)\n self.up6 = UNetUp(256, 64)\n\n self.final = nn.Sequential(nn.ConvTranspose2d(128, channels, 4, stride=2, padding=1), nn.Tanh())\n\n def forward(self, x):\n # Propogate noise through fc layer and reshape to img shape\n d1 = self.down1(x)\n d2 = self.down2(d1)\n d3 = self.down3(d2)\n d4 = self.down4(d3)\n d5 = self.down5(d4)\n d6 = self.down6(d5)\n d7 = self.down7(d6)\n u1 = self.up1(d7, d6)\n u2 = self.up2(u1, d5)\n u3 = self.up3(u2, d4)\n u4 = self.up4(u3, d3)\n u5 = self.up5(u4, d2)\n u6 = self.up6(u5, d1)\n\n return self.final(u6)\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, in_channels=3):\n super(Discriminator, self).__init__()\n\n def discrimintor_block(in_features, out_features, normalize=True):\n \"\"\"Discriminator block\"\"\"\n layers = [nn.Conv2d(in_features, out_features, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.BatchNorm2d(out_features, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *discrimintor_block(in_channels, 64, normalize=False),\n *discrimintor_block(64, 128),\n *discrimintor_block(128, 256),\n nn.ZeroPad2d((1, 0, 1, 0)),\n nn.Conv2d(256, 1, kernel_size=4)\n )\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/ebgan/ebgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=62, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"number of image channels\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise):\n out = self.l1(noise)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n # Upsampling\n self.down = nn.Sequential(nn.Conv2d(opt.channels, 64, 3, 2, 1), nn.ReLU())\n # Fully-connected layers\n self.down_size = opt.img_size // 2\n down_dim = 64 * (opt.img_size // 2) ** 2\n\n self.embedding = nn.Linear(down_dim, 32)\n\n self.fc = nn.Sequential(\n nn.BatchNorm1d(32, 0.8),\n nn.ReLU(inplace=True),\n nn.Linear(32, down_dim),\n nn.BatchNorm1d(down_dim),\n nn.ReLU(inplace=True),\n )\n # Upsampling\n self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv2d(64, opt.channels, 3, 1, 1))\n\n def forward(self, img):\n out = self.down(img)\n embedding = self.embedding(out.view(out.size(0), -1))\n out = self.fc(embedding)\n out = self.up(out.view(out.size(0), 64, self.down_size, self.down_size))\n return out, embedding\n\n\n# Reconstruction loss of AE\npixelwise_loss = nn.MSELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n pixelwise_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef pullaway_loss(embeddings):\n norm = torch.sqrt(torch.sum(embeddings ** 2, -1, keepdim=True))\n normalized_emb = embeddings / norm\n similarity = torch.matmul(normalized_emb, normalized_emb.transpose(1, 0))\n batch_size = embeddings.size(0)\n loss_pt = (torch.sum(similarity) - batch_size) / (batch_size * (batch_size - 1))\n return loss_pt\n\n\n# ----------\n# Training\n# ----------\n\n# BEGAN hyper parameters\nlambda_pt = 0.1\nmargin = max(1, opt.batch_size / 64.0)\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n recon_imgs, img_embeddings = discriminator(gen_imgs)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = pixelwise_loss(recon_imgs, gen_imgs.detach()) + lambda_pt * pullaway_loss(img_embeddings)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_recon, _ = discriminator(real_imgs)\n fake_recon, _ = discriminator(gen_imgs.detach())\n\n d_loss_real = pixelwise_loss(real_recon, real_imgs)\n d_loss_fake = pixelwise_loss(fake_recon, gen_imgs.detach())\n\n d_loss = d_loss_real\n if (margin - d_loss_fake.data).item() > 0:\n d_loss += margin - d_loss_fake\n\n d_loss.backward()\n optimizer_D.step()\n\n # --------------\n # Log Progress\n # --------------\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/esrgan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n# Normalization parameters for pre-trained PyTorch models\nmean = np.array([0.485, 0.456, 0.406])\nstd = np.array([0.229, 0.224, 0.225])\n\n\ndef denormalize(tensors):\n \"\"\" Denormalizes image tensors using mean and std \"\"\"\n for c in range(3):\n tensors[:, c].mul_(std[c]).add_(mean[c])\n return torch.clamp(tensors, 0, 255)\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, hr_shape):\n hr_height, hr_width = hr_shape\n # Transforms for low resolution images and high resolution images\n self.lr_transform = transforms.Compose(\n [\n transforms.Resize((hr_height // 4, hr_height // 4), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize(mean, std),\n ]\n )\n self.hr_transform = transforms.Compose(\n [\n transforms.Resize((hr_height, hr_height), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize(mean, std),\n ]\n )\n\n self.files = sorted(glob.glob(root + \"/*.*\"))\n\n def __getitem__(self, index):\n img = Image.open(self.files[index % len(self.files)])\n img_lr = self.lr_transform(img)\n img_hr = self.hr_transform(img)\n\n return {\"lr\": img_lr, \"hr\": img_hr}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/esrgan/esrgan.py",
"content": "\"\"\"\nSuper-resolution of CelebA using Generative Adversarial Networks.\nThe dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0\n(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)\nInstrustion on running the script:\n1. Download the dataset from the provided link\n2. Save the folder 'img_align_celeba' to '../../data/'\n4. Run the sript using command 'python3 esrgan.py'\n\"\"\"\n\nimport argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image, make_grid\n\nfrom torch.utils.data import DataLoader\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images/training\", exist_ok=True)\nos.makedirs(\"saved_models\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"img_align_celeba\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=4, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.9, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--hr_height\", type=int, default=256, help=\"high res. image height\")\nparser.add_argument(\"--hr_width\", type=int, default=256, help=\"high res. image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving image samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=5000, help=\"batch interval between model checkpoints\")\nparser.add_argument(\"--residual_blocks\", type=int, default=23, help=\"number of residual blocks in the generator\")\nparser.add_argument(\"--warmup_batches\", type=int, default=500, help=\"number of batches with pixel-wise loss only\")\nparser.add_argument(\"--lambda_adv\", type=float, default=5e-3, help=\"adversarial loss weight\")\nparser.add_argument(\"--lambda_pixel\", type=float, default=1e-2, help=\"pixel-wise loss weight\")\nopt = parser.parse_args()\nprint(opt)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nhr_shape = (opt.hr_height, opt.hr_width)\n\n# Initialize generator and discriminator\ngenerator = GeneratorRRDB(opt.channels, filters=64, num_res_blocks=opt.residual_blocks).to(device)\ndiscriminator = Discriminator(input_shape=(opt.channels, *hr_shape)).to(device)\nfeature_extractor = FeatureExtractor().to(device)\n\n# Set feature extractor to inference mode\nfeature_extractor.eval()\n\n# Losses\ncriterion_GAN = torch.nn.BCEWithLogitsLoss().to(device)\ncriterion_content = torch.nn.L1Loss().to(device)\ncriterion_pixel = torch.nn.L1Loss().to(device)\n\nif opt.epoch != 0:\n # Load pretrained models\n generator.load_state_dict(torch.load(\"saved_models/generator_%d.pth\" % opt.epoch))\n discriminator.load_state_dict(torch.load(\"saved_models/discriminator_%d.pth\" % opt.epoch))\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor\n\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, hr_shape=hr_shape),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, imgs in enumerate(dataloader):\n\n batches_done = epoch * len(dataloader) + i\n\n # Configure model input\n imgs_lr = Variable(imgs[\"lr\"].type(Tensor))\n imgs_hr = Variable(imgs[\"hr\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n\n # ------------------\n # Train Generators\n # ------------------\n\n optimizer_G.zero_grad()\n\n # Generate a high resolution image from low resolution input\n gen_hr = generator(imgs_lr)\n\n # Measure pixel-wise loss against ground truth\n loss_pixel = criterion_pixel(gen_hr, imgs_hr)\n\n if batches_done < opt.warmup_batches:\n # Warm-up (pixel-wise loss only)\n loss_pixel.backward()\n optimizer_G.step()\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [G pixel: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), loss_pixel.item())\n )\n continue\n\n # Extract validity predictions from discriminator\n pred_real = discriminator(imgs_hr).detach()\n pred_fake = discriminator(gen_hr)\n\n # Adversarial loss (relativistic average GAN)\n loss_GAN = criterion_GAN(pred_fake - pred_real.mean(0, keepdim=True), valid)\n\n # Content loss\n gen_features = feature_extractor(gen_hr)\n real_features = feature_extractor(imgs_hr).detach()\n loss_content = criterion_content(gen_features, real_features)\n\n # Total generator loss\n loss_G = loss_content + opt.lambda_adv * loss_GAN + opt.lambda_pixel * loss_pixel\n\n loss_G.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n pred_real = discriminator(imgs_hr)\n pred_fake = discriminator(gen_hr.detach())\n\n # Adversarial loss for real and fake images (relativistic average GAN)\n loss_real = criterion_GAN(pred_real - pred_fake.mean(0, keepdim=True), valid)\n loss_fake = criterion_GAN(pred_fake - pred_real.mean(0, keepdim=True), fake)\n\n # Total loss\n loss_D = (loss_real + loss_fake) / 2\n\n loss_D.backward()\n optimizer_D.step()\n\n # --------------\n # Log Progress\n # --------------\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, content: %f, adv: %f, pixel: %f]\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D.item(),\n loss_G.item(),\n loss_content.item(),\n loss_GAN.item(),\n loss_pixel.item(),\n )\n )\n\n if batches_done % opt.sample_interval == 0:\n # Save image grid with upsampled inputs and ESRGAN outputs\n imgs_lr = nn.functional.interpolate(imgs_lr, scale_factor=4)\n img_grid = denormalize(torch.cat((imgs_lr, gen_hr), -1))\n save_image(img_grid, \"images/training/%d.png\" % batches_done, nrow=1, normalize=False)\n\n if batches_done % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(generator.state_dict(), \"saved_models/generator_%d.pth\" % epoch)\n torch.save(discriminator.state_dict(), \"saved_models/discriminator_%d.pth\" %epoch)\n"
},
{
"path": "implementations/esrgan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torchvision.models import vgg19\nimport math\n\n\nclass FeatureExtractor(nn.Module):\n def __init__(self):\n super(FeatureExtractor, self).__init__()\n vgg19_model = vgg19(pretrained=True)\n self.vgg19_54 = nn.Sequential(*list(vgg19_model.features.children())[:35])\n\n def forward(self, img):\n return self.vgg19_54(img)\n\n\nclass DenseResidualBlock(nn.Module):\n \"\"\"\n The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)\n \"\"\"\n\n def __init__(self, filters, res_scale=0.2):\n super(DenseResidualBlock, self).__init__()\n self.res_scale = res_scale\n\n def block(in_features, non_linearity=True):\n layers = [nn.Conv2d(in_features, filters, 3, 1, 1, bias=True)]\n if non_linearity:\n layers += [nn.LeakyReLU()]\n return nn.Sequential(*layers)\n\n self.b1 = block(in_features=1 * filters)\n self.b2 = block(in_features=2 * filters)\n self.b3 = block(in_features=3 * filters)\n self.b4 = block(in_features=4 * filters)\n self.b5 = block(in_features=5 * filters, non_linearity=False)\n self.blocks = [self.b1, self.b2, self.b3, self.b4, self.b5]\n\n def forward(self, x):\n inputs = x\n for block in self.blocks:\n out = block(inputs)\n inputs = torch.cat([inputs, out], 1)\n return out.mul(self.res_scale) + x\n\n\nclass ResidualInResidualDenseBlock(nn.Module):\n def __init__(self, filters, res_scale=0.2):\n super(ResidualInResidualDenseBlock, self).__init__()\n self.res_scale = res_scale\n self.dense_blocks = nn.Sequential(\n DenseResidualBlock(filters), DenseResidualBlock(filters), DenseResidualBlock(filters)\n )\n\n def forward(self, x):\n return self.dense_blocks(x).mul(self.res_scale) + x\n\n\nclass GeneratorRRDB(nn.Module):\n def __init__(self, channels, filters=64, num_res_blocks=16, num_upsample=2):\n super(GeneratorRRDB, self).__init__()\n\n # First layer\n self.conv1 = nn.Conv2d(channels, filters, kernel_size=3, stride=1, padding=1)\n # Residual blocks\n self.res_blocks = nn.Sequential(*[ResidualInResidualDenseBlock(filters) for _ in range(num_res_blocks)])\n # Second conv layer post residual blocks\n self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1)\n # Upsampling layers\n upsample_layers = []\n for _ in range(num_upsample):\n upsample_layers += [\n nn.Conv2d(filters, filters * 4, kernel_size=3, stride=1, padding=1),\n nn.LeakyReLU(),\n nn.PixelShuffle(upscale_factor=2),\n ]\n self.upsampling = nn.Sequential(*upsample_layers)\n # Final output block\n self.conv3 = nn.Sequential(\n nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1),\n nn.LeakyReLU(),\n nn.Conv2d(filters, channels, kernel_size=3, stride=1, padding=1),\n )\n\n def forward(self, x):\n out1 = self.conv1(x)\n out = self.res_blocks(out1)\n out2 = self.conv2(out)\n out = torch.add(out1, out2)\n out = self.upsampling(out)\n out = self.conv3(out)\n return out\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n\n self.input_shape = input_shape\n in_channels, in_height, in_width = self.input_shape\n patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)\n self.output_shape = (1, patch_h, patch_w)\n\n def discriminator_block(in_filters, out_filters, first_block=False):\n layers = []\n layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1))\n if not first_block:\n layers.append(nn.BatchNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n layers.append(nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=2, padding=1))\n layers.append(nn.BatchNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n layers = []\n in_filters = in_channels\n for i, out_filters in enumerate([64, 128, 256, 512]):\n layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))\n in_filters = out_filters\n\n layers.append(nn.Conv2d(out_filters, 1, kernel_size=3, stride=1, padding=1))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/esrgan/test_on_image.py",
"content": "from models import GeneratorRRDB\nfrom datasets import denormalize, mean, std\nimport torch\nfrom torch.autograd import Variable\nimport argparse\nimport os\nfrom torchvision import transforms\nfrom torchvision.utils import save_image\nfrom PIL import Image\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--image_path\", type=str, required=True, help=\"Path to image\")\nparser.add_argument(\"--checkpoint_model\", type=str, required=True, help=\"Path to checkpoint model\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"Number of image channels\")\nparser.add_argument(\"--residual_blocks\", type=int, default=23, help=\"Number of residual blocks in G\")\nopt = parser.parse_args()\nprint(opt)\n\nos.makedirs(\"images/outputs\", exist_ok=True)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Define model and load model checkpoint\ngenerator = GeneratorRRDB(opt.channels, filters=64, num_res_blocks=opt.residual_blocks).to(device)\ngenerator.load_state_dict(torch.load(opt.checkpoint_model))\ngenerator.eval()\n\ntransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])\n\n# Prepare input\nimage_tensor = Variable(transform(Image.open(opt.image_path))).to(device).unsqueeze(0)\n\n# Upsample image\nwith torch.no_grad():\n sr_image = denormalize(generator(image_tensor)).cpu()\n\n# Save image\nfn = opt.image_path.split(\"/\")[-1]\nsave_image(sr_image, f\"images/outputs/sr-{fn}\")\n"
},
{
"path": "implementations/gan/gan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.size(0), *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n nn.Sigmoid(),\n )\n\n def forward(self, img):\n img_flat = img.view(img.size(0), -1)\n validity = self.model(img_flat)\n\n return validity\n\n\n# Loss function\nadversarial_loss = torch.nn.BCELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.size(0), 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.size(0), 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(real_imgs), valid)\n fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/infogan/infogan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images/static/\", exist_ok=True)\nos.makedirs(\"images/varying_c1/\", exist_ok=True)\nos.makedirs(\"images/varying_c2/\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=62, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--code_dim\", type=int, default=2, help=\"latent code\")\nparser.add_argument(\"--n_classes\", type=int, default=10, help=\"number of classes for dataset\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\ndef to_categorical(y, num_columns):\n \"\"\"Returns one-hot encoded Variable\"\"\"\n y_cat = np.zeros((y.shape[0], num_columns))\n y_cat[range(y.shape[0]), y] = 1.0\n\n return Variable(FloatTensor(y_cat))\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n input_dim = opt.latent_dim + opt.n_classes + opt.code_dim\n\n self.init_size = opt.img_size // 4 # Initial size before upsampling\n self.l1 = nn.Sequential(nn.Linear(input_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise, labels, code):\n gen_input = torch.cat((noise, labels, code), -1)\n out = self.l1(gen_input)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n \"\"\"Returns layers of each discriminator block\"\"\"\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.conv_blocks = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n\n # Output layers\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1))\n self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.n_classes), nn.Softmax())\n self.latent_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.code_dim))\n\n def forward(self, img):\n out = self.conv_blocks(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n label = self.aux_layer(out)\n latent_code = self.latent_layer(out)\n\n return validity, label, latent_code\n\n\n# Loss functions\nadversarial_loss = torch.nn.MSELoss()\ncategorical_loss = torch.nn.CrossEntropyLoss()\ncontinuous_loss = torch.nn.MSELoss()\n\n# Loss weights\nlambda_cat = 1\nlambda_con = 0.1\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n categorical_loss.cuda()\n continuous_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_info = torch.optim.Adam(\n itertools.chain(generator.parameters(), discriminator.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n# Static generator inputs for sampling\nstatic_z = Variable(FloatTensor(np.zeros((opt.n_classes ** 2, opt.latent_dim))))\nstatic_label = to_categorical(\n np.array([num for _ in range(opt.n_classes) for num in range(opt.n_classes)]), num_columns=opt.n_classes\n)\nstatic_code = Variable(FloatTensor(np.zeros((opt.n_classes ** 2, opt.code_dim))))\n\n\ndef sample_image(n_row, batches_done):\n \"\"\"Saves a grid of generated digits ranging from 0 to n_classes\"\"\"\n # Static sample\n z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))\n static_sample = generator(z, static_label, static_code)\n save_image(static_sample.data, \"images/static/%d.png\" % batches_done, nrow=n_row, normalize=True)\n\n # Get varied c1 and c2\n zeros = np.zeros((n_row ** 2, 1))\n c_varied = np.repeat(np.linspace(-1, 1, n_row)[:, np.newaxis], n_row, 0)\n c1 = Variable(FloatTensor(np.concatenate((c_varied, zeros), -1)))\n c2 = Variable(FloatTensor(np.concatenate((zeros, c_varied), -1)))\n sample1 = generator(static_z, static_label, c1)\n sample2 = generator(static_z, static_label, c2)\n save_image(sample1.data, \"images/varying_c1/%d.png\" % batches_done, nrow=n_row, normalize=True)\n save_image(sample2.data, \"images/varying_c2/%d.png\" % batches_done, nrow=n_row, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, labels) in enumerate(dataloader):\n\n batch_size = imgs.shape[0]\n\n # Adversarial ground truths\n valid = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)\n fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(FloatTensor))\n labels = to_categorical(labels.numpy(), num_columns=opt.n_classes)\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise and labels as generator input\n z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n label_input = to_categorical(np.random.randint(0, opt.n_classes, batch_size), num_columns=opt.n_classes)\n code_input = Variable(FloatTensor(np.random.uniform(-1, 1, (batch_size, opt.code_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z, label_input, code_input)\n\n # Loss measures generator's ability to fool the discriminator\n validity, _, _ = discriminator(gen_imgs)\n g_loss = adversarial_loss(validity, valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Loss for real images\n real_pred, _, _ = discriminator(real_imgs)\n d_real_loss = adversarial_loss(real_pred, valid)\n\n # Loss for fake images\n fake_pred, _, _ = discriminator(gen_imgs.detach())\n d_fake_loss = adversarial_loss(fake_pred, fake)\n\n # Total discriminator loss\n d_loss = (d_real_loss + d_fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n # ------------------\n # Information Loss\n # ------------------\n\n optimizer_info.zero_grad()\n\n # Sample labels\n sampled_labels = np.random.randint(0, opt.n_classes, batch_size)\n\n # Ground truth labels\n gt_labels = Variable(LongTensor(sampled_labels), requires_grad=False)\n\n # Sample noise, labels and code as generator input\n z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n label_input = to_categorical(sampled_labels, num_columns=opt.n_classes)\n code_input = Variable(FloatTensor(np.random.uniform(-1, 1, (batch_size, opt.code_dim))))\n\n gen_imgs = generator(z, label_input, code_input)\n _, pred_label, pred_code = discriminator(gen_imgs)\n\n info_loss = lambda_cat * categorical_loss(pred_label, gt_labels) + lambda_con * continuous_loss(\n pred_code, code_input\n )\n\n info_loss.backward()\n optimizer_info.step()\n\n # --------------\n # Log Progress\n # --------------\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [info loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item(), info_loss.item())\n )\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n sample_image(n_row=10, batches_done=batches_done)\n"
},
{
"path": "implementations/lsgan/lsgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=1000, help=\"number of image channels\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, z):\n out = self.l1(z)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.model = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n self.adv_layer = nn.Linear(128 * ds_size ** 2, 1)\n\n def forward(self, img):\n out = self.model(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n\n return validity\n\n\n# !!! Minimizes MSE instead of BCE\nadversarial_loss = torch.nn.MSELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(real_imgs), valid)\n fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n d_loss = 0.5 * (real_loss + fake_loss)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/munit/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n\n self.files = sorted(glob.glob(os.path.join(root, mode) + \"/*.*\"))\n if mode == \"train\":\n self.files.extend(sorted(glob.glob(os.path.join(root, \"test\") + \"/*.*\")))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n w, h = img.size\n img_A = img.crop((0, 0, w / 2, h))\n img_B = img.crop((w / 2, 0, w, h))\n\n if np.random.random() < 0.5:\n img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\")\n img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], \"RGB\")\n\n img_A = self.transform(img_A)\n img_B = self.transform(img_B)\n\n return {\"A\": img_A, \"B\": img_B}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/munit/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torch.autograd import Variable\nimport numpy as np\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass LambdaLR:\n def __init__(self, n_epochs, offset, decay_start_epoch):\n assert (n_epochs - decay_start_epoch) > 0, \"Decay must start before the training session ends!\"\n self.n_epochs = n_epochs\n self.offset = offset\n self.decay_start_epoch = decay_start_epoch\n\n def step(self, epoch):\n return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)\n\n\n#################################\n# Encoder\n#################################\n\n\nclass Encoder(nn.Module):\n def __init__(self, in_channels=3, dim=64, n_residual=3, n_downsample=2, style_dim=8):\n super(Encoder, self).__init__()\n self.content_encoder = ContentEncoder(in_channels, dim, n_residual, n_downsample)\n self.style_encoder = StyleEncoder(in_channels, dim, n_downsample, style_dim)\n\n def forward(self, x):\n content_code = self.content_encoder(x)\n style_code = self.style_encoder(x)\n return content_code, style_code\n\n\n#################################\n# Decoder\n#################################\n\n\nclass Decoder(nn.Module):\n def __init__(self, out_channels=3, dim=64, n_residual=3, n_upsample=2, style_dim=8):\n super(Decoder, self).__init__()\n\n layers = []\n dim = dim * 2 ** n_upsample\n # Residual blocks\n for _ in range(n_residual):\n layers += [ResidualBlock(dim, norm=\"adain\")]\n\n # Upsampling\n for _ in range(n_upsample):\n layers += [\n nn.Upsample(scale_factor=2),\n nn.Conv2d(dim, dim // 2, 5, stride=1, padding=2),\n LayerNorm(dim // 2),\n nn.ReLU(inplace=True),\n ]\n dim = dim // 2\n\n # Output layer\n layers += [nn.ReflectionPad2d(3), nn.Conv2d(dim, out_channels, 7), nn.Tanh()]\n\n self.model = nn.Sequential(*layers)\n\n # Initiate mlp (predicts AdaIN parameters)\n num_adain_params = self.get_num_adain_params()\n self.mlp = MLP(style_dim, num_adain_params)\n\n def get_num_adain_params(self):\n \"\"\"Return the number of AdaIN parameters needed by the model\"\"\"\n num_adain_params = 0\n for m in self.modules():\n if m.__class__.__name__ == \"AdaptiveInstanceNorm2d\":\n num_adain_params += 2 * m.num_features\n return num_adain_params\n\n def assign_adain_params(self, adain_params):\n \"\"\"Assign the adain_params to the AdaIN layers in model\"\"\"\n for m in self.modules():\n if m.__class__.__name__ == \"AdaptiveInstanceNorm2d\":\n # Extract mean and std predictions\n mean = adain_params[:, : m.num_features]\n std = adain_params[:, m.num_features : 2 * m.num_features]\n # Update bias and weight\n m.bias = mean.contiguous().view(-1)\n m.weight = std.contiguous().view(-1)\n # Move pointer\n if adain_params.size(1) > 2 * m.num_features:\n adain_params = adain_params[:, 2 * m.num_features :]\n\n def forward(self, content_code, style_code):\n # Update AdaIN parameters by MLP prediction based off style code\n self.assign_adain_params(self.mlp(style_code))\n img = self.model(content_code)\n return img\n\n\n#################################\n# Content Encoder\n#################################\n\n\nclass ContentEncoder(nn.Module):\n def __init__(self, in_channels=3, dim=64, n_residual=3, n_downsample=2):\n super(ContentEncoder, self).__init__()\n\n # Initial convolution block\n layers = [\n nn.ReflectionPad2d(3),\n nn.Conv2d(in_channels, dim, 7),\n nn.InstanceNorm2d(dim),\n nn.ReLU(inplace=True),\n ]\n\n # Downsampling\n for _ in range(n_downsample):\n layers += [\n nn.Conv2d(dim, dim * 2, 4, stride=2, padding=1),\n nn.InstanceNorm2d(dim * 2),\n nn.ReLU(inplace=True),\n ]\n dim *= 2\n\n # Residual blocks\n for _ in range(n_residual):\n layers += [ResidualBlock(dim, norm=\"in\")]\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\n#################################\n# Style Encoder\n#################################\n\n\nclass StyleEncoder(nn.Module):\n def __init__(self, in_channels=3, dim=64, n_downsample=2, style_dim=8):\n super(StyleEncoder, self).__init__()\n\n # Initial conv block\n layers = [nn.ReflectionPad2d(3), nn.Conv2d(in_channels, dim, 7), nn.ReLU(inplace=True)]\n\n # Downsampling\n for _ in range(2):\n layers += [nn.Conv2d(dim, dim * 2, 4, stride=2, padding=1), nn.ReLU(inplace=True)]\n dim *= 2\n\n # Downsampling with constant depth\n for _ in range(n_downsample - 2):\n layers += [nn.Conv2d(dim, dim, 4, stride=2, padding=1), nn.ReLU(inplace=True)]\n\n # Average pool and output layer\n layers += [nn.AdaptiveAvgPool2d(1), nn.Conv2d(dim, style_dim, 1, 1, 0)]\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\n######################################\n# MLP (predicts AdaIn parameters)\n######################################\n\n\nclass MLP(nn.Module):\n def __init__(self, input_dim, output_dim, dim=256, n_blk=3, activ=\"relu\"):\n super(MLP, self).__init__()\n layers = [nn.Linear(input_dim, dim), nn.ReLU(inplace=True)]\n for _ in range(n_blk - 2):\n layers += [nn.Linear(dim, dim), nn.ReLU(inplace=True)]\n layers += [nn.Linear(dim, output_dim)]\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x.view(x.size(0), -1))\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass MultiDiscriminator(nn.Module):\n def __init__(self, in_channels=3):\n super(MultiDiscriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, normalize=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n # Extracts three discriminator models\n self.models = nn.ModuleList()\n for i in range(3):\n self.models.add_module(\n \"disc_%d\" % i,\n nn.Sequential(\n *discriminator_block(in_channels, 64, normalize=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n *discriminator_block(256, 512),\n nn.Conv2d(512, 1, 3, padding=1)\n ),\n )\n\n self.downsample = nn.AvgPool2d(in_channels, stride=2, padding=[1, 1], count_include_pad=False)\n\n def compute_loss(self, x, gt):\n \"\"\"Computes the MSE between model output and scalar gt\"\"\"\n loss = sum([torch.mean((out - gt) ** 2) for out in self.forward(x)])\n return loss\n\n def forward(self, x):\n outputs = []\n for m in self.models:\n outputs.append(m(x))\n x = self.downsample(x)\n return outputs\n\n\n##############################\n# Custom Blocks\n##############################\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, features, norm=\"in\"):\n super(ResidualBlock, self).__init__()\n\n norm_layer = AdaptiveInstanceNorm2d if norm == \"adain\" else nn.InstanceNorm2d\n\n self.block = nn.Sequential(\n nn.ReflectionPad2d(1),\n nn.Conv2d(features, features, 3),\n norm_layer(features),\n nn.ReLU(inplace=True),\n nn.ReflectionPad2d(1),\n nn.Conv2d(features, features, 3),\n norm_layer(features),\n )\n\n def forward(self, x):\n return x + self.block(x)\n\n\n##############################\n# Custom Layers\n##############################\n\n\nclass AdaptiveInstanceNorm2d(nn.Module):\n \"\"\"Reference: https://github.com/NVlabs/MUNIT/blob/master/networks.py\"\"\"\n\n def __init__(self, num_features, eps=1e-5, momentum=0.1):\n super(AdaptiveInstanceNorm2d, self).__init__()\n self.num_features = num_features\n self.eps = eps\n self.momentum = momentum\n # weight and bias are dynamically assigned\n self.weight = None\n self.bias = None\n # just dummy buffers, not used\n self.register_buffer(\"running_mean\", torch.zeros(num_features))\n self.register_buffer(\"running_var\", torch.ones(num_features))\n\n def forward(self, x):\n assert (\n self.weight is not None and self.bias is not None\n ), \"Please assign weight and bias before calling AdaIN!\"\n b, c, h, w = x.size()\n running_mean = self.running_mean.repeat(b)\n running_var = self.running_var.repeat(b)\n\n # Apply instance norm\n x_reshaped = x.contiguous().view(1, b * c, h, w)\n\n out = F.batch_norm(\n x_reshaped, running_mean, running_var, self.weight, self.bias, True, self.momentum, self.eps\n )\n\n return out.view(b, c, h, w)\n\n def __repr__(self):\n return self.__class__.__name__ + \"(\" + str(self.num_features) + \")\"\n\n\nclass LayerNorm(nn.Module):\n def __init__(self, num_features, eps=1e-5, affine=True):\n super(LayerNorm, self).__init__()\n self.num_features = num_features\n self.affine = affine\n self.eps = eps\n\n if self.affine:\n self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_())\n self.beta = nn.Parameter(torch.zeros(num_features))\n\n def forward(self, x):\n shape = [-1] + [1] * (x.dim() - 1)\n mean = x.view(x.size(0), -1).mean(1).view(*shape)\n std = x.view(x.size(0), -1).std(1).view(*shape)\n x = (x - mean) / (std + self.eps)\n\n if self.affine:\n shape = [1, -1] + [1] * (x.dim() - 2)\n x = x * self.gamma.view(*shape) + self.beta.view(*shape)\n return x\n"
},
{
"path": "implementations/munit/munit.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport datetime\nimport time\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"edges2shoes\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=1, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0001, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=128, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=128, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval saving generator samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between saving model checkpoints\")\nparser.add_argument(\"--n_downsample\", type=int, default=2, help=\"number downsampling layers in encoder\")\nparser.add_argument(\"--n_residual\", type=int, default=3, help=\"number of residual blocks in encoder / decoder\")\nparser.add_argument(\"--dim\", type=int, default=64, help=\"number of filters in first encoder layer\")\nparser.add_argument(\"--style_dim\", type=int, default=8, help=\"dimensionality of the style code\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = torch.cuda.is_available()\n\n# Create sample and checkpoint directories\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\ncriterion_recon = torch.nn.L1Loss()\n\n# Initialize encoders, generators and discriminators\nEnc1 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)\nDec1 = Decoder(dim=opt.dim, n_upsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)\nEnc2 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)\nDec2 = Decoder(dim=opt.dim, n_upsample=opt.n_downsample, n_residual=opt.n_residual, style_dim=opt.style_dim)\nD1 = MultiDiscriminator()\nD2 = MultiDiscriminator()\n\nif cuda:\n Enc1 = Enc1.cuda()\n Dec1 = Dec1.cuda()\n Enc2 = Enc2.cuda()\n Dec2 = Dec2.cuda()\n D1 = D1.cuda()\n D2 = D2.cuda()\n criterion_recon.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n Enc1.load_state_dict(torch.load(\"saved_models/%s/Enc1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n Dec1.load_state_dict(torch.load(\"saved_models/%s/Dec1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n Enc2.load_state_dict(torch.load(\"saved_models/%s/Enc2_%d.pth\" % (opt.dataset_name, opt.epoch)))\n Dec2.load_state_dict(torch.load(\"saved_models/%s/Dec2_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D1.load_state_dict(torch.load(\"saved_models/%s/D1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D2.load_state_dict(torch.load(\"saved_models/%s/D2_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n Enc1.apply(weights_init_normal)\n Dec1.apply(weights_init_normal)\n Enc2.apply(weights_init_normal)\n Dec2.apply(weights_init_normal)\n D1.apply(weights_init_normal)\n D2.apply(weights_init_normal)\n\n# Loss weights\nlambda_gan = 1\nlambda_id = 10\nlambda_style = 1\nlambda_cont = 1\nlambda_cyc = 0\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(Enc1.parameters(), Dec1.parameters(), Enc2.parameters(), Dec2.parameters()),\n lr=opt.lr,\n betas=(opt.b1, opt.b2),\n)\noptimizer_D1 = torch.optim.Adam(D1.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D2 = torch.optim.Adam(D2.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Learning rate update schedulers\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n optimizer_G, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D1 = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D1, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D2 = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D2, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\n# Configure dataloaders\ntransforms_ = [\n transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, mode=\"val\"),\n batch_size=5,\n shuffle=True,\n num_workers=1,\n)\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the validation set\"\"\"\n imgs = next(iter(val_dataloader))\n img_samples = None\n for img1, img2 in zip(imgs[\"A\"], imgs[\"B\"]):\n # Create copies of image\n X1 = img1.unsqueeze(0).repeat(opt.style_dim, 1, 1, 1)\n X1 = Variable(X1.type(Tensor))\n # Get random style codes\n s_code = np.random.uniform(-1, 1, (opt.style_dim, opt.style_dim))\n s_code = Variable(Tensor(s_code))\n # Generate samples\n c_code_1, _ = Enc1(X1)\n X12 = Dec2(c_code_1, s_code)\n # Concatenate samples horisontally\n X12 = torch.cat([x for x in X12.data.cpu()], -1)\n img_sample = torch.cat((img1, X12), -1).unsqueeze(0)\n # Concatenate with previous samples vertically\n img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)\n save_image(img_samples, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=5, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\n# Adversarial ground truths\nvalid = 1\nfake = 0\n\nprev_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Set model input\n X1 = Variable(batch[\"A\"].type(Tensor))\n X2 = Variable(batch[\"B\"].type(Tensor))\n\n # Sampled style codes\n style_1 = Variable(torch.randn(X1.size(0), opt.style_dim, 1, 1).type(Tensor))\n style_2 = Variable(torch.randn(X1.size(0), opt.style_dim, 1, 1).type(Tensor))\n\n # -------------------------------\n # Train Encoders and Generators\n # -------------------------------\n\n optimizer_G.zero_grad()\n\n # Get shared latent representation\n c_code_1, s_code_1 = Enc1(X1)\n c_code_2, s_code_2 = Enc2(X2)\n\n # Reconstruct images\n X11 = Dec1(c_code_1, s_code_1)\n X22 = Dec2(c_code_2, s_code_2)\n\n # Translate images\n X21 = Dec1(c_code_2, style_1)\n X12 = Dec2(c_code_1, style_2)\n\n # Cycle translation\n c_code_21, s_code_21 = Enc1(X21)\n c_code_12, s_code_12 = Enc2(X12)\n X121 = Dec1(c_code_12, s_code_1) if lambda_cyc > 0 else 0\n X212 = Dec2(c_code_21, s_code_2) if lambda_cyc > 0 else 0\n\n # Losses\n loss_GAN_1 = lambda_gan * D1.compute_loss(X21, valid)\n loss_GAN_2 = lambda_gan * D2.compute_loss(X12, valid)\n loss_ID_1 = lambda_id * criterion_recon(X11, X1)\n loss_ID_2 = lambda_id * criterion_recon(X22, X2)\n loss_s_1 = lambda_style * criterion_recon(s_code_21, style_1)\n loss_s_2 = lambda_style * criterion_recon(s_code_12, style_2)\n loss_c_1 = lambda_cont * criterion_recon(c_code_12, c_code_1.detach())\n loss_c_2 = lambda_cont * criterion_recon(c_code_21, c_code_2.detach())\n loss_cyc_1 = lambda_cyc * criterion_recon(X121, X1) if lambda_cyc > 0 else 0\n loss_cyc_2 = lambda_cyc * criterion_recon(X212, X2) if lambda_cyc > 0 else 0\n\n # Total loss\n loss_G = (\n loss_GAN_1\n + loss_GAN_2\n + loss_ID_1\n + loss_ID_2\n + loss_s_1\n + loss_s_2\n + loss_c_1\n + loss_c_2\n + loss_cyc_1\n + loss_cyc_2\n )\n\n loss_G.backward()\n optimizer_G.step()\n\n # -----------------------\n # Train Discriminator 1\n # -----------------------\n\n optimizer_D1.zero_grad()\n\n loss_D1 = D1.compute_loss(X1, valid) + D1.compute_loss(X21.detach(), fake)\n\n loss_D1.backward()\n optimizer_D1.step()\n\n # -----------------------\n # Train Discriminator 2\n # -----------------------\n\n optimizer_D2.zero_grad()\n\n loss_D2 = D2.compute_loss(X2, valid) + D2.compute_loss(X12.detach(), fake)\n\n loss_D2.backward()\n optimizer_D2.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] ETA: %s\"\n % (epoch, opt.n_epochs, i, len(dataloader), (loss_D1 + loss_D2).item(), loss_G.item(), time_left)\n )\n\n # If at sample interval save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n # Update learning rates\n lr_scheduler_G.step()\n lr_scheduler_D1.step()\n lr_scheduler_D2.step()\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(Enc1.state_dict(), \"saved_models/%s/Enc1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(Dec1.state_dict(), \"saved_models/%s/Dec1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(Enc2.state_dict(), \"saved_models/%s/Enc2_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(Dec2.state_dict(), \"saved_models/%s/Dec2_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D1.state_dict(), \"saved_models/%s/D1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D2.state_dict(), \"saved_models/%s/D2_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/pix2pix/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n\n self.files = sorted(glob.glob(os.path.join(root, mode) + \"/*.*\"))\n if mode == \"train\":\n self.files.extend(sorted(glob.glob(os.path.join(root, \"test\") + \"/*.*\")))\n\n def __getitem__(self, index):\n\n img = Image.open(self.files[index % len(self.files)])\n w, h = img.size\n img_A = img.crop((0, 0, w / 2, h))\n img_B = img.crop((w / 2, 0, w, h))\n\n if np.random.random() < 0.5:\n img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\")\n img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], \"RGB\")\n\n img_A = self.transform(img_A)\n img_B = self.transform(img_B)\n\n return {\"A\": img_A, \"B\": img_B}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/pix2pix/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\n##############################\n# U-NET\n##############################\n\n\nclass UNetDown(nn.Module):\n def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n super(UNetDown, self).__init__()\n layers = [nn.Conv2d(in_size, out_size, 4, 2, 1, bias=False)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_size))\n layers.append(nn.LeakyReLU(0.2))\n if dropout:\n layers.append(nn.Dropout(dropout))\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass UNetUp(nn.Module):\n def __init__(self, in_size, out_size, dropout=0.0):\n super(UNetUp, self).__init__()\n layers = [\n nn.ConvTranspose2d(in_size, out_size, 4, 2, 1, bias=False),\n nn.InstanceNorm2d(out_size),\n nn.ReLU(inplace=True),\n ]\n if dropout:\n layers.append(nn.Dropout(dropout))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x, skip_input):\n x = self.model(x)\n x = torch.cat((x, skip_input), 1)\n\n return x\n\n\nclass GeneratorUNet(nn.Module):\n def __init__(self, in_channels=3, out_channels=3):\n super(GeneratorUNet, self).__init__()\n\n self.down1 = UNetDown(in_channels, 64, normalize=False)\n self.down2 = UNetDown(64, 128)\n self.down3 = UNetDown(128, 256)\n self.down4 = UNetDown(256, 512, dropout=0.5)\n self.down5 = UNetDown(512, 512, dropout=0.5)\n self.down6 = UNetDown(512, 512, dropout=0.5)\n self.down7 = UNetDown(512, 512, dropout=0.5)\n self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)\n\n self.up1 = UNetUp(512, 512, dropout=0.5)\n self.up2 = UNetUp(1024, 512, dropout=0.5)\n self.up3 = UNetUp(1024, 512, dropout=0.5)\n self.up4 = UNetUp(1024, 512, dropout=0.5)\n self.up5 = UNetUp(1024, 256)\n self.up6 = UNetUp(512, 128)\n self.up7 = UNetUp(256, 64)\n\n self.final = nn.Sequential(\n nn.Upsample(scale_factor=2),\n nn.ZeroPad2d((1, 0, 1, 0)),\n nn.Conv2d(128, out_channels, 4, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, x):\n # U-Net generator with skip connections from encoder to decoder\n d1 = self.down1(x)\n d2 = self.down2(d1)\n d3 = self.down3(d2)\n d4 = self.down4(d3)\n d5 = self.down5(d4)\n d6 = self.down6(d5)\n d7 = self.down7(d6)\n d8 = self.down8(d7)\n u1 = self.up1(d8, d7)\n u2 = self.up2(u1, d6)\n u3 = self.up3(u2, d5)\n u4 = self.up4(u3, d4)\n u5 = self.up5(u4, d3)\n u6 = self.up6(u5, d2)\n u7 = self.up7(u6, d1)\n\n return self.final(u7)\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, in_channels=3):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, normalization=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalization:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *discriminator_block(in_channels * 2, 64, normalization=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n *discriminator_block(256, 512),\n nn.ZeroPad2d((1, 0, 1, 0)),\n nn.Conv2d(512, 1, 4, padding=1, bias=False)\n )\n\n def forward(self, img_A, img_B):\n # Concatenate image and condition image by channels to produce input\n img_input = torch.cat((img_A, img_B), 1)\n return self.model(img_input)\n"
},
{
"path": "implementations/pix2pix/pix2pix.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport time\nimport datetime\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"facades\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=1, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=256, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=256, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\n \"--sample_interval\", type=int, default=500, help=\"interval between sampling of images from generators\"\n)\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nopt = parser.parse_args()\nprint(opt)\n\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\ncuda = True if torch.cuda.is_available() else False\n\n# Loss functions\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_pixelwise = torch.nn.L1Loss()\n\n# Loss weight of L1 pixel-wise loss between translated image and real image\nlambda_pixel = 100\n\n# Calculate output of image discriminator (PatchGAN)\npatch = (1, opt.img_height // 2 ** 4, opt.img_width // 2 ** 4)\n\n# Initialize generator and discriminator\ngenerator = GeneratorUNet()\ndiscriminator = Discriminator()\n\nif cuda:\n generator = generator.cuda()\n discriminator = discriminator.cuda()\n criterion_GAN.cuda()\n criterion_pixelwise.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n generator.load_state_dict(torch.load(\"saved_models/%s/generator_%d.pth\" % (opt.dataset_name, opt.epoch)))\n discriminator.load_state_dict(torch.load(\"saved_models/%s/discriminator_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n generator.apply(weights_init_normal)\n discriminator.apply(weights_init_normal)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Configure dataloaders\ntransforms_ = [\n transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, mode=\"val\"),\n batch_size=10,\n shuffle=True,\n num_workers=1,\n)\n\n# Tensor type\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the validation set\"\"\"\n imgs = next(iter(val_dataloader))\n real_A = Variable(imgs[\"B\"].type(Tensor))\n real_B = Variable(imgs[\"A\"].type(Tensor))\n fake_B = generator(real_A)\n img_sample = torch.cat((real_A.data, fake_B.data, real_B.data), -2)\n save_image(img_sample, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=5, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nprev_time = time.time()\n\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Model inputs\n real_A = Variable(batch[\"B\"].type(Tensor))\n real_B = Variable(batch[\"A\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((real_A.size(0), *patch))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((real_A.size(0), *patch))), requires_grad=False)\n\n # ------------------\n # Train Generators\n # ------------------\n\n optimizer_G.zero_grad()\n\n # GAN loss\n fake_B = generator(real_A)\n pred_fake = discriminator(fake_B, real_A)\n loss_GAN = criterion_GAN(pred_fake, valid)\n # Pixel-wise loss\n loss_pixel = criterion_pixelwise(fake_B, real_B)\n\n # Total loss\n loss_G = loss_GAN + lambda_pixel * loss_pixel\n\n loss_G.backward()\n\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Real loss\n pred_real = discriminator(real_B, real_A)\n loss_real = criterion_GAN(pred_real, valid)\n\n # Fake loss\n pred_fake = discriminator(fake_B.detach(), real_A)\n loss_fake = criterion_GAN(pred_fake, fake)\n\n # Total loss\n loss_D = 0.5 * (loss_real + loss_fake)\n\n loss_D.backward()\n optimizer_D.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, pixel: %f, adv: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D.item(),\n loss_G.item(),\n loss_pixel.item(),\n loss_GAN.item(),\n time_left,\n )\n )\n\n # If at sample interval save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(generator.state_dict(), \"saved_models/%s/generator_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(discriminator.state_dict(), \"saved_models/%s/discriminator_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/pixelda/mnistm.py",
"content": "\"\"\"Dataset setting and data loader for MNIST-M.\n\nModified from\nhttps://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py\n\nCREDIT: https://github.com/corenel\n\"\"\"\n\nfrom __future__ import print_function\n\nimport errno\nimport os\n\nimport torch\nimport torch.utils.data as data\nfrom PIL import Image\n\n\nclass MNISTM(data.Dataset):\n \"\"\"`MNIST-M Dataset.\"\"\"\n\n url = \"https://github.com/VanushVaswani/keras_mnistm/releases/download/1.0/keras_mnistm.pkl.gz\"\n\n raw_folder = 'raw'\n processed_folder = 'processed'\n training_file = 'mnist_m_train.pt'\n test_file = 'mnist_m_test.pt'\n\n def __init__(self,\n root, mnist_root=\"data\",\n train=True,\n transform=None, target_transform=None,\n download=False):\n \"\"\"Init MNIST-M dataset.\"\"\"\n super(MNISTM, self).__init__()\n self.root = os.path.expanduser(root)\n self.mnist_root = os.path.expanduser(mnist_root)\n self.transform = transform\n self.target_transform = target_transform\n self.train = train # training set or test set\n\n if download:\n self.download()\n\n if not self._check_exists():\n raise RuntimeError('Dataset not found.' +\n ' You can use download=True to download it')\n\n if self.train:\n self.train_data, self.train_labels = \\\n torch.load(os.path.join(self.root,\n self.processed_folder,\n self.training_file))\n else:\n self.test_data, self.test_labels = \\\n torch.load(os.path.join(self.root,\n self.processed_folder,\n self.test_file))\n\n def __getitem__(self, index):\n \"\"\"Get images and target for data loader.\n\n Args:\n index (int): Index\n\n Returns:\n tuple: (image, target) where target is index of the target class.\n \"\"\"\n if self.train:\n img, target = self.train_data[index], self.train_labels[index]\n else:\n img, target = self.test_data[index], self.test_labels[index]\n\n # doing this so that it is consistent with all other datasets\n # to return a PIL Image\n img = Image.fromarray(img.squeeze().numpy(), mode='RGB')\n\n if self.transform is not None:\n img = self.transform(img)\n\n if self.target_transform is not None:\n target = self.target_transform(target)\n\n return img, target\n\n def __len__(self):\n \"\"\"Return size of dataset.\"\"\"\n if self.train:\n return len(self.train_data)\n else:\n return len(self.test_data)\n\n def _check_exists(self):\n return os.path.exists(os.path.join(self.root,\n self.processed_folder,\n self.training_file)) and \\\n os.path.exists(os.path.join(self.root,\n self.processed_folder,\n self.test_file))\n\n def download(self):\n \"\"\"Download the MNIST data.\"\"\"\n # import essential packages\n from six.moves import urllib\n import gzip\n import pickle\n from torchvision import datasets\n\n # check if dataset already exists\n if self._check_exists():\n return\n\n # make data dirs\n try:\n os.makedirs(os.path.join(self.root, self.raw_folder))\n os.makedirs(os.path.join(self.root, self.processed_folder))\n except OSError as e:\n if e.errno == errno.EEXIST:\n pass\n else:\n raise\n\n # download pkl files\n print('Downloading ' + self.url)\n filename = self.url.rpartition('/')[2]\n file_path = os.path.join(self.root, self.raw_folder, filename)\n if not os.path.exists(file_path.replace('.gz', '')):\n data = urllib.request.urlopen(self.url)\n with open(file_path, 'wb') as f:\n f.write(data.read())\n with open(file_path.replace('.gz', ''), 'wb') as out_f, \\\n gzip.GzipFile(file_path) as zip_f:\n out_f.write(zip_f.read())\n os.unlink(file_path)\n\n # process and save as torch files\n print('Processing...')\n\n # load MNIST-M images from pkl file\n with open(file_path.replace('.gz', ''), \"rb\") as f:\n mnist_m_data = pickle.load(f, encoding='bytes')\n mnist_m_train_data = torch.ByteTensor(mnist_m_data[b'train'])\n mnist_m_test_data = torch.ByteTensor(mnist_m_data[b'test'])\n\n # get MNIST labels\n mnist_train_labels = datasets.MNIST(root=self.mnist_root,\n train=True,\n download=True).train_labels\n mnist_test_labels = datasets.MNIST(root=self.mnist_root,\n train=False,\n download=True).test_labels\n\n # save MNIST-M dataset\n training_set = (mnist_m_train_data, mnist_train_labels)\n test_set = (mnist_m_test_data, mnist_test_labels)\n with open(os.path.join(self.root,\n self.processed_folder,\n self.training_file), 'wb') as f:\n torch.save(training_set, f)\n with open(os.path.join(self.root,\n self.processed_folder,\n self.test_file), 'wb') as f:\n torch.save(test_set, f)\n\n print('Done!')\n"
},
{
"path": "implementations/pixelda/pixelda.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom mnistm import MNISTM\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--n_residual_blocks\", type=int, default=6, help=\"number of residual blocks in generator\")\nparser.add_argument(\"--latent_dim\", type=int, default=10, help=\"dimensionality of the noise input\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--n_classes\", type=int, default=10, help=\"number of classes in the dataset\")\nparser.add_argument(\"--sample_interval\", type=int, default=300, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\n# Calculate output of image discriminator (PatchGAN)\npatch = int(opt.img_size / 2 ** 4)\npatch = (1, patch, patch)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, in_features=64, out_features=64):\n super(ResidualBlock, self).__init__()\n\n self.block = nn.Sequential(\n nn.Conv2d(in_features, in_features, 3, 1, 1),\n nn.BatchNorm2d(in_features),\n nn.ReLU(inplace=True),\n nn.Conv2d(in_features, in_features, 3, 1, 1),\n nn.BatchNorm2d(in_features),\n )\n\n def forward(self, x):\n return x + self.block(x)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n # Fully-connected layer which constructs image channel shaped output from noise\n self.fc = nn.Linear(opt.latent_dim, opt.channels * opt.img_size ** 2)\n\n self.l1 = nn.Sequential(nn.Conv2d(opt.channels * 2, 64, 3, 1, 1), nn.ReLU(inplace=True))\n\n resblocks = []\n for _ in range(opt.n_residual_blocks):\n resblocks.append(ResidualBlock())\n self.resblocks = nn.Sequential(*resblocks)\n\n self.l2 = nn.Sequential(nn.Conv2d(64, opt.channels, 3, 1, 1), nn.Tanh())\n\n def forward(self, img, z):\n gen_input = torch.cat((img, self.fc(z).view(*img.shape)), 1)\n out = self.l1(gen_input)\n out = self.resblocks(out)\n img_ = self.l2(out)\n\n return img_\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def block(in_features, out_features, normalization=True):\n \"\"\"Discriminator block\"\"\"\n layers = [nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True)]\n if normalization:\n layers.append(nn.InstanceNorm2d(out_features))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.channels, 64, normalization=False),\n *block(64, 128),\n *block(128, 256),\n *block(256, 512),\n nn.Conv2d(512, 1, 3, 1, 1)\n )\n\n def forward(self, img):\n validity = self.model(img)\n\n return validity\n\n\nclass Classifier(nn.Module):\n def __init__(self):\n super(Classifier, self).__init__()\n\n def block(in_features, out_features, normalization=True):\n \"\"\"Classifier block\"\"\"\n layers = [nn.Conv2d(in_features, out_features, 3, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True)]\n if normalization:\n layers.append(nn.InstanceNorm2d(out_features))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.channels, 64, normalization=False), *block(64, 128), *block(128, 256), *block(256, 512)\n )\n\n input_size = opt.img_size // 2 ** 4\n self.output_layer = nn.Sequential(nn.Linear(512 * input_size ** 2, opt.n_classes), nn.Softmax())\n\n def forward(self, img):\n feature_repr = self.model(img)\n feature_repr = feature_repr.view(feature_repr.size(0), -1)\n label = self.output_layer(feature_repr)\n return label\n\n\n# Loss function\nadversarial_loss = torch.nn.MSELoss()\ntask_loss = torch.nn.CrossEntropyLoss()\n\n# Loss weights\nlambda_adv = 1\nlambda_task = 0.1\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\nclassifier = Classifier()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n classifier.cuda()\n adversarial_loss.cuda()\n task_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\nclassifier.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader_A = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\nos.makedirs(\"../../data/mnistm\", exist_ok=True)\ndataloader_B = torch.utils.data.DataLoader(\n MNISTM(\n \"../../data/mnistm\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [\n transforms.Resize(opt.img_size),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n ]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(generator.parameters(), classifier.parameters()), lr=opt.lr, betas=(opt.b1, opt.b2)\n)\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n# ----------\n# Training\n# ----------\n\n# Keeps 100 accuracy measurements\ntask_performance = []\ntarget_performance = []\n\nfor epoch in range(opt.n_epochs):\n for i, ((imgs_A, labels_A), (imgs_B, labels_B)) in enumerate(zip(dataloader_A, dataloader_B)):\n\n batch_size = imgs_A.size(0)\n\n # Adversarial ground truths\n valid = Variable(FloatTensor(batch_size, *patch).fill_(1.0), requires_grad=False)\n fake = Variable(FloatTensor(batch_size, *patch).fill_(0.0), requires_grad=False)\n\n # Configure input\n imgs_A = Variable(imgs_A.type(FloatTensor).expand(batch_size, 3, opt.img_size, opt.img_size))\n labels_A = Variable(labels_A.type(LongTensor))\n imgs_B = Variable(imgs_B.type(FloatTensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise\n z = Variable(FloatTensor(np.random.uniform(-1, 1, (batch_size, opt.latent_dim))))\n\n # Generate a batch of images\n fake_B = generator(imgs_A, z)\n\n # Perform task on translated source image\n label_pred = classifier(fake_B)\n\n # Calculate the task loss\n task_loss_ = (task_loss(label_pred, labels_A) + task_loss(classifier(imgs_A), labels_A)) / 2\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = lambda_adv * adversarial_loss(discriminator(fake_B), valid) + lambda_task * task_loss_\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Measure discriminator's ability to classify real from generated samples\n real_loss = adversarial_loss(discriminator(imgs_B), valid)\n fake_loss = adversarial_loss(discriminator(fake_B.detach()), fake)\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n # ---------------------------------------\n # Evaluate Performance on target domain\n # ---------------------------------------\n\n # Evaluate performance on translated Domain A\n acc = np.mean(np.argmax(label_pred.data.cpu().numpy(), axis=1) == labels_A.data.cpu().numpy())\n task_performance.append(acc)\n if len(task_performance) > 100:\n task_performance.pop(0)\n\n # Evaluate performance on Domain B\n pred_B = classifier(imgs_B)\n target_acc = np.mean(np.argmax(pred_B.data.cpu().numpy(), axis=1) == labels_B.numpy())\n target_performance.append(target_acc)\n if len(target_performance) > 100:\n target_performance.pop(0)\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] [CLF acc: %3d%% (%3d%%), target_acc: %3d%% (%3d%%)]\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader_A),\n d_loss.item(),\n g_loss.item(),\n 100 * acc,\n 100 * np.mean(task_performance),\n 100 * target_acc,\n 100 * np.mean(target_performance),\n )\n )\n\n batches_done = len(dataloader_A) * epoch + i\n if batches_done % opt.sample_interval == 0:\n sample = torch.cat((imgs_A.data[:5], fake_B.data[:5], imgs_B.data[:5]), -2)\n save_image(sample, \"images/%d.png\" % batches_done, nrow=int(math.sqrt(batch_size)), normalize=True)\n"
},
{
"path": "implementations/relativistic_gan/relativistic_gan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nparser.add_argument(\"--rel_avg_gan\", action=\"store_true\", help=\"relativistic average GAN instead of standard\")\nopt = parser.parse_args()\nprint(opt)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.init_size = opt.img_size // 4\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, z):\n out = self.l1(z)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.model = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1))\n\n def forward(self, img):\n out = self.model(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n\n return validity\n\n\n# Loss function\nadversarial_loss = torch.nn.BCEWithLogitsLoss().to(device)\n\n# Initialize generator and discriminator\ngenerator = Generator().to(device)\ndiscriminator = Discriminator().to(device)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Adversarial ground truths\n valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n real_pred = discriminator(real_imgs).detach()\n fake_pred = discriminator(gen_imgs)\n\n if opt.rel_avg_gan:\n g_loss = adversarial_loss(fake_pred - real_pred.mean(0, keepdim=True), valid)\n else:\n g_loss = adversarial_loss(fake_pred - real_pred, valid)\n\n # Loss measures generator's ability to fool the discriminator\n g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Predict validity\n real_pred = discriminator(real_imgs)\n fake_pred = discriminator(gen_imgs.detach())\n\n if opt.rel_avg_gan:\n real_loss = adversarial_loss(real_pred - fake_pred.mean(0, keepdim=True), valid)\n fake_loss = adversarial_loss(fake_pred - real_pred.mean(0, keepdim=True), fake)\n else:\n real_loss = adversarial_loss(real_pred - fake_pred, valid)\n fake_loss = adversarial_loss(fake_pred - real_pred, fake)\n\n d_loss = (real_loss + fake_loss) / 2\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/sgan/sgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--num_classes\", type=int, default=10, help=\"number of classes for dataset\")\nparser.add_argument(\"--img_size\", type=int, default=32, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between image sampling\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n self.label_emb = nn.Embedding(opt.num_classes, opt.latent_dim)\n\n self.init_size = opt.img_size // 4 # Initial size before upsampling\n self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))\n\n self.conv_blocks = nn.Sequential(\n nn.BatchNorm2d(128),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 128, 3, stride=1, padding=1),\n nn.BatchNorm2d(128, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Upsample(scale_factor=2),\n nn.Conv2d(128, 64, 3, stride=1, padding=1),\n nn.BatchNorm2d(64, 0.8),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n nn.Tanh(),\n )\n\n def forward(self, noise):\n out = self.l1(noise)\n out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n img = self.conv_blocks(out)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n def discriminator_block(in_filters, out_filters, bn=True):\n \"\"\"Returns layers of each discriminator block\"\"\"\n block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]\n if bn:\n block.append(nn.BatchNorm2d(out_filters, 0.8))\n return block\n\n self.conv_blocks = nn.Sequential(\n *discriminator_block(opt.channels, 16, bn=False),\n *discriminator_block(16, 32),\n *discriminator_block(32, 64),\n *discriminator_block(64, 128),\n )\n\n # The height and width of downsampled image\n ds_size = opt.img_size // 2 ** 4\n\n # Output layers\n self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())\n self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.num_classes + 1), nn.Softmax())\n\n def forward(self, img):\n out = self.conv_blocks(img)\n out = out.view(out.shape[0], -1)\n validity = self.adv_layer(out)\n label = self.aux_layer(out)\n\n return validity, label\n\n\n# Loss functions\nadversarial_loss = torch.nn.BCELoss()\nauxiliary_loss = torch.nn.CrossEntropyLoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n auxiliary_loss.cuda()\n\n# Initialize weights\ngenerator.apply(weights_init_normal)\ndiscriminator.apply(weights_init_normal)\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nFloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\nLongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, labels) in enumerate(dataloader):\n\n batch_size = imgs.shape[0]\n\n # Adversarial ground truths\n valid = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)\n fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)\n fake_aux_gt = Variable(LongTensor(batch_size).fill_(opt.num_classes), requires_grad=False)\n\n # Configure input\n real_imgs = Variable(imgs.type(FloatTensor))\n labels = Variable(labels.type(LongTensor))\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Sample noise and labels as generator input\n z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))\n\n # Generate a batch of images\n gen_imgs = generator(z)\n\n # Loss measures generator's ability to fool the discriminator\n validity, _ = discriminator(gen_imgs)\n g_loss = adversarial_loss(validity, valid)\n\n g_loss.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Loss for real images\n real_pred, real_aux = discriminator(real_imgs)\n d_real_loss = (adversarial_loss(real_pred, valid) + auxiliary_loss(real_aux, labels)) / 2\n\n # Loss for fake images\n fake_pred, fake_aux = discriminator(gen_imgs.detach())\n d_fake_loss = (adversarial_loss(fake_pred, fake) + auxiliary_loss(fake_aux, fake_aux_gt)) / 2\n\n # Total discriminator loss\n d_loss = (d_real_loss + d_fake_loss) / 2\n\n # Calculate discriminator accuracy\n pred = np.concatenate([real_aux.data.cpu().numpy(), fake_aux.data.cpu().numpy()], axis=0)\n gt = np.concatenate([labels.data.cpu().numpy(), fake_aux_gt.data.cpu().numpy()], axis=0)\n d_acc = np.mean(np.argmax(pred, axis=1) == gt)\n\n d_loss.backward()\n optimizer_D.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), 100 * d_acc, g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/softmax_gan/softmax_gan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(opt.img_size ** 2, 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n )\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n validity = self.model(img_flat)\n\n return validity\n\n\n# Loss function\nadversarial_loss = torch.nn.BCELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n adversarial_loss.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef log(x):\n return torch.log(x + 1e-8)\n\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n optimizer_G.zero_grad()\n optimizer_D.zero_grad()\n\n batch_size = imgs.shape[0]\n\n # Adversarial ground truths\n g_target = 1 / (batch_size * 2)\n d_target = 1 / batch_size\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n # Generate a batch of images\n gen_imgs = generator(z)\n\n d_real = discriminator(real_imgs)\n d_fake = discriminator(gen_imgs)\n\n # Partition function\n Z = torch.sum(torch.exp(-d_real)) + torch.sum(torch.exp(-d_fake))\n\n # Calculate loss of discriminator and update\n d_loss = d_target * torch.sum(d_real) + log(Z)\n d_loss.backward(retain_graph=True)\n optimizer_D.step()\n\n # Calculate loss of generator and update\n g_loss = g_target * (torch.sum(d_real) + torch.sum(d_fake)) + log(Z)\n g_loss.backward()\n optimizer_G.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n"
},
{
"path": "implementations/srgan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n# Normalization parameters for pre-trained PyTorch models\nmean = np.array([0.485, 0.456, 0.406])\nstd = np.array([0.229, 0.224, 0.225])\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, hr_shape):\n hr_height, hr_width = hr_shape\n # Transforms for low resolution images and high resolution images\n self.lr_transform = transforms.Compose(\n [\n transforms.Resize((hr_height // 4, hr_height // 4), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize(mean, std),\n ]\n )\n self.hr_transform = transforms.Compose(\n [\n transforms.Resize((hr_height, hr_height), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize(mean, std),\n ]\n )\n\n self.files = sorted(glob.glob(root + \"/*.*\"))\n\n def __getitem__(self, index):\n img = Image.open(self.files[index % len(self.files)])\n img_lr = self.lr_transform(img)\n img_hr = self.hr_transform(img)\n\n return {\"lr\": img_lr, \"hr\": img_hr}\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/srgan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torchvision.models import vgg19\nimport math\n\n\nclass FeatureExtractor(nn.Module):\n def __init__(self):\n super(FeatureExtractor, self).__init__()\n vgg19_model = vgg19(pretrained=True)\n self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:18])\n\n def forward(self, img):\n return self.feature_extractor(img)\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, in_features):\n super(ResidualBlock, self).__init__()\n self.conv_block = nn.Sequential(\n nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(in_features, 0.8),\n nn.PReLU(),\n nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(in_features, 0.8),\n )\n\n def forward(self, x):\n return x + self.conv_block(x)\n\n\nclass GeneratorResNet(nn.Module):\n def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16):\n super(GeneratorResNet, self).__init__()\n\n # First layer\n self.conv1 = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=9, stride=1, padding=4), nn.PReLU())\n\n # Residual blocks\n res_blocks = []\n for _ in range(n_residual_blocks):\n res_blocks.append(ResidualBlock(64))\n self.res_blocks = nn.Sequential(*res_blocks)\n\n # Second conv layer post residual blocks\n self.conv2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64, 0.8))\n\n # Upsampling layers\n upsampling = []\n for out_features in range(2):\n upsampling += [\n # nn.Upsample(scale_factor=2),\n nn.Conv2d(64, 256, 3, 1, 1),\n nn.BatchNorm2d(256),\n nn.PixelShuffle(upscale_factor=2),\n nn.PReLU(),\n ]\n self.upsampling = nn.Sequential(*upsampling)\n\n # Final output layer\n self.conv3 = nn.Sequential(nn.Conv2d(64, out_channels, kernel_size=9, stride=1, padding=4), nn.Tanh())\n\n def forward(self, x):\n out1 = self.conv1(x)\n out = self.res_blocks(out1)\n out2 = self.conv2(out)\n out = torch.add(out1, out2)\n out = self.upsampling(out)\n out = self.conv3(out)\n return out\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n\n self.input_shape = input_shape\n in_channels, in_height, in_width = self.input_shape\n patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)\n self.output_shape = (1, patch_h, patch_w)\n\n def discriminator_block(in_filters, out_filters, first_block=False):\n layers = []\n layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1))\n if not first_block:\n layers.append(nn.BatchNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n layers.append(nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=2, padding=1))\n layers.append(nn.BatchNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n layers = []\n in_filters = in_channels\n for i, out_filters in enumerate([64, 128, 256, 512]):\n layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))\n in_filters = out_filters\n\n layers.append(nn.Conv2d(out_filters, 1, kernel_size=3, stride=1, padding=1))\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/srgan/srgan.py",
"content": "\"\"\"\nSuper-resolution of CelebA using Generative Adversarial Networks.\nThe dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0\n(if not available there see if options are listed at http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)\nInstrustion on running the script:\n1. Download the dataset from the provided link\n2. Save the folder 'img_align_celeba' to '../../data/'\n4. Run the sript using command 'python3 srgan.py'\n\"\"\"\n\nimport argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image, make_grid\n\nfrom torch.utils.data import DataLoader\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\nos.makedirs(\"saved_models\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"img_align_celeba\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=4, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--hr_height\", type=int, default=256, help=\"high res. image height\")\nparser.add_argument(\"--hr_width\", type=int, default=256, help=\"high res. image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving image samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = torch.cuda.is_available()\n\nhr_shape = (opt.hr_height, opt.hr_width)\n\n# Initialize generator and discriminator\ngenerator = GeneratorResNet()\ndiscriminator = Discriminator(input_shape=(opt.channels, *hr_shape))\nfeature_extractor = FeatureExtractor()\n\n# Set feature extractor to inference mode\nfeature_extractor.eval()\n\n# Losses\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_content = torch.nn.L1Loss()\n\nif cuda:\n generator = generator.cuda()\n discriminator = discriminator.cuda()\n feature_extractor = feature_extractor.cuda()\n criterion_GAN = criterion_GAN.cuda()\n criterion_content = criterion_content.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n generator.load_state_dict(torch.load(\"saved_models/generator_%d.pth\"))\n discriminator.load_state_dict(torch.load(\"saved_models/discriminator_%d.pth\"))\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, hr_shape=hr_shape),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\n# ----------\n# Training\n# ----------\n\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, imgs in enumerate(dataloader):\n\n # Configure model input\n imgs_lr = Variable(imgs[\"lr\"].type(Tensor))\n imgs_hr = Variable(imgs[\"hr\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n\n # ------------------\n # Train Generators\n # ------------------\n\n optimizer_G.zero_grad()\n\n # Generate a high resolution image from low resolution input\n gen_hr = generator(imgs_lr)\n\n # Adversarial loss\n loss_GAN = criterion_GAN(discriminator(gen_hr), valid)\n\n # Content loss\n gen_features = feature_extractor(gen_hr)\n real_features = feature_extractor(imgs_hr)\n loss_content = criterion_content(gen_features, real_features.detach())\n\n # Total loss\n loss_G = loss_content + 1e-3 * loss_GAN\n\n loss_G.backward()\n optimizer_G.step()\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Loss of real and fake images\n loss_real = criterion_GAN(discriminator(imgs_hr), valid)\n loss_fake = criterion_GAN(discriminator(gen_hr.detach()), fake)\n\n # Total loss\n loss_D = (loss_real + loss_fake) / 2\n\n loss_D.backward()\n optimizer_D.step()\n\n # --------------\n # Log Progress\n # --------------\n\n sys.stdout.write(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), loss_D.item(), loss_G.item())\n )\n\n batches_done = epoch * len(dataloader) + i\n if batches_done % opt.sample_interval == 0:\n # Save image grid with upsampled inputs and SRGAN outputs\n imgs_lr = nn.functional.interpolate(imgs_lr, scale_factor=4)\n gen_hr = make_grid(gen_hr, nrow=1, normalize=True)\n imgs_lr = make_grid(imgs_lr, nrow=1, normalize=True)\n img_grid = torch.cat((imgs_lr, gen_hr), -1)\n save_image(img_grid, \"images/%d.png\" % batches_done, normalize=False)\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(generator.state_dict(), \"saved_models/generator_%d.pth\" % epoch)\n torch.save(discriminator.state_dict(), \"saved_models/discriminator_%d.pth\" % epoch)\n"
},
{
"path": "implementations/stargan/datasets.py",
"content": "import glob\nimport random\nimport os\nimport numpy as np\nimport torch\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass CelebADataset(Dataset):\n def __init__(self, root, transforms_=None, mode=\"train\", attributes=None):\n self.transform = transforms.Compose(transforms_)\n\n self.selected_attrs = attributes\n self.files = sorted(glob.glob(\"%s/*.jpg\" % root))\n self.files = self.files[:-2000] if mode == \"train\" else self.files[-2000:]\n self.label_path = glob.glob(\"%s/*.txt\" % root)[0]\n self.annotations = self.get_annotations()\n\n def get_annotations(self):\n \"\"\"Extracts annotations for CelebA\"\"\"\n annotations = {}\n lines = [line.rstrip() for line in open(self.label_path, \"r\")]\n self.label_names = lines[1].split()\n for _, line in enumerate(lines[2:]):\n filename, *values = line.split()\n labels = []\n for attr in self.selected_attrs:\n idx = self.label_names.index(attr)\n labels.append(1 * (values[idx] == \"1\"))\n annotations[filename] = labels\n return annotations\n\n def __getitem__(self, index):\n filepath = self.files[index % len(self.files)]\n filename = filepath.split(\"/\")[-1]\n img = self.transform(Image.open(filepath))\n label = self.annotations[filename]\n label = torch.FloatTensor(np.array(label))\n\n return img, label\n\n def __len__(self):\n return len(self.files)\n"
},
{
"path": "implementations/stargan/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n\n\n##############################\n# RESNET\n##############################\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, in_features):\n super(ResidualBlock, self).__init__()\n\n conv_block = [\n nn.Conv2d(in_features, in_features, 3, stride=1, padding=1, bias=False),\n nn.InstanceNorm2d(in_features, affine=True, track_running_stats=True),\n nn.ReLU(inplace=True),\n nn.Conv2d(in_features, in_features, 3, stride=1, padding=1, bias=False),\n nn.InstanceNorm2d(in_features, affine=True, track_running_stats=True),\n ]\n\n self.conv_block = nn.Sequential(*conv_block)\n\n def forward(self, x):\n return x + self.conv_block(x)\n\n\nclass GeneratorResNet(nn.Module):\n def __init__(self, img_shape=(3, 128, 128), res_blocks=9, c_dim=5):\n super(GeneratorResNet, self).__init__()\n channels, img_size, _ = img_shape\n\n # Initial convolution block\n model = [\n nn.Conv2d(channels + c_dim, 64, 7, stride=1, padding=3, bias=False),\n nn.InstanceNorm2d(64, affine=True, track_running_stats=True),\n nn.ReLU(inplace=True),\n ]\n\n # Downsampling\n curr_dim = 64\n for _ in range(2):\n model += [\n nn.Conv2d(curr_dim, curr_dim * 2, 4, stride=2, padding=1, bias=False),\n nn.InstanceNorm2d(curr_dim * 2, affine=True, track_running_stats=True),\n nn.ReLU(inplace=True),\n ]\n curr_dim *= 2\n\n # Residual blocks\n for _ in range(res_blocks):\n model += [ResidualBlock(curr_dim)]\n\n # Upsampling\n for _ in range(2):\n model += [\n nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, stride=2, padding=1, bias=False),\n nn.InstanceNorm2d(curr_dim // 2, affine=True, track_running_stats=True),\n nn.ReLU(inplace=True),\n ]\n curr_dim = curr_dim // 2\n\n # Output layer\n model += [nn.Conv2d(curr_dim, channels, 7, stride=1, padding=3), nn.Tanh()]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x, c):\n c = c.view(c.size(0), c.size(1), 1, 1)\n c = c.repeat(1, 1, x.size(2), x.size(3))\n x = torch.cat((x, c), 1)\n return self.model(x)\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, img_shape=(3, 128, 128), c_dim=5, n_strided=6):\n super(Discriminator, self).__init__()\n channels, img_size, _ = img_shape\n\n def discriminator_block(in_filters, out_filters):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1), nn.LeakyReLU(0.01)]\n return layers\n\n layers = discriminator_block(channels, 64)\n curr_dim = 64\n for _ in range(n_strided - 1):\n layers.extend(discriminator_block(curr_dim, curr_dim * 2))\n curr_dim *= 2\n\n self.model = nn.Sequential(*layers)\n\n # Output 1: PatchGAN\n self.out1 = nn.Conv2d(curr_dim, 1, 3, padding=1, bias=False)\n # Output 2: Class prediction\n kernel_size = img_size // 2 ** n_strided\n self.out2 = nn.Conv2d(curr_dim, c_dim, kernel_size, bias=False)\n\n def forward(self, img):\n feature_repr = self.model(img)\n out_adv = self.out1(feature_repr)\n out_cls = self.out2(feature_repr)\n return out_adv, out_cls.view(out_cls.size(0), -1)\n"
},
{
"path": "implementations/stargan/stargan.py",
"content": "\"\"\"\nStarGAN (CelebA)\nThe dataset can be downloaded from: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=0\nAnd the annotations: https://www.dropbox.com/sh/8oqt9vytwxb3s4r/AAA8YmAHNNU6BEfWMPMfM6r9a/Anno?dl=0&preview=list_attr_celeba.txt\nInstructions on running the script:\n1. Download the dataset and annotations from the provided link\n2. Copy 'list_attr_celeba.txt' to folder 'img_align_celeba'\n2. Save the folder 'img_align_celeba' to '../../data/'\n4. Run the script by 'python3 stargan.py'\n\"\"\"\n\n\nimport argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport time\nimport datetime\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\nimport torch.autograd as autograd\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\nos.makedirs(\"saved_models\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"img_align_celeba\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=16, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=128, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=128, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval between saving generator samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between model checkpoints\")\nparser.add_argument(\"--residual_blocks\", type=int, default=6, help=\"number of residual blocks in generator\")\nparser.add_argument(\n \"--selected_attrs\",\n \"--list\",\n nargs=\"+\",\n help=\"selected attributes for the CelebA dataset\",\n default=[\"Black_Hair\", \"Blond_Hair\", \"Brown_Hair\", \"Male\", \"Young\"],\n)\nparser.add_argument(\"--n_critic\", type=int, default=5, help=\"number of training iterations for WGAN discriminator\")\nopt = parser.parse_args()\nprint(opt)\n\nc_dim = len(opt.selected_attrs)\nimg_shape = (opt.channels, opt.img_height, opt.img_width)\n\ncuda = torch.cuda.is_available()\n\n# Loss functions\ncriterion_cycle = torch.nn.L1Loss()\n\n\ndef criterion_cls(logit, target):\n return F.binary_cross_entropy_with_logits(logit, target, size_average=False) / logit.size(0)\n\n\n# Loss weights\nlambda_cls = 1\nlambda_rec = 10\nlambda_gp = 10\n\n# Initialize generator and discriminator\ngenerator = GeneratorResNet(img_shape=img_shape, res_blocks=opt.residual_blocks, c_dim=c_dim)\ndiscriminator = Discriminator(img_shape=img_shape, c_dim=c_dim)\n\nif cuda:\n generator = generator.cuda()\n discriminator = discriminator.cuda()\n criterion_cycle.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n generator.load_state_dict(torch.load(\"saved_models/generator_%d.pth\" % opt.epoch))\n discriminator.load_state_dict(torch.load(\"saved_models/discriminator_%d.pth\" % opt.epoch))\nelse:\n generator.apply(weights_init_normal)\n discriminator.apply(weights_init_normal)\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Configure dataloaders\ntrain_transforms = [\n transforms.Resize(int(1.12 * opt.img_height), Image.BICUBIC),\n transforms.RandomCrop(opt.img_height),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\ndataloader = DataLoader(\n CelebADataset(\n \"../../data/%s\" % opt.dataset_name, transforms_=train_transforms, mode=\"train\", attributes=opt.selected_attrs\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n\nval_transforms = [\n transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\nval_dataloader = DataLoader(\n CelebADataset(\n \"../../data/%s\" % opt.dataset_name, transforms_=val_transforms, mode=\"val\", attributes=opt.selected_attrs\n ),\n batch_size=10,\n shuffle=True,\n num_workers=1,\n)\n\n# Tensor type\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef compute_gradient_penalty(D, real_samples, fake_samples):\n \"\"\"Calculates the gradient penalty loss for WGAN GP\"\"\"\n # Random weight term for interpolation between real and fake samples\n alpha = Tensor(np.random.random((real_samples.size(0), 1, 1, 1)))\n # Get random interpolation between real and fake samples\n interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)\n d_interpolates, _ = D(interpolates)\n fake = Variable(Tensor(np.ones(d_interpolates.shape)), requires_grad=False)\n # Get gradient w.r.t. interpolates\n gradients = autograd.grad(\n outputs=d_interpolates,\n inputs=interpolates,\n grad_outputs=fake,\n create_graph=True,\n retain_graph=True,\n only_inputs=True,\n )[0]\n gradients = gradients.view(gradients.size(0), -1)\n gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n return gradient_penalty\n\n\nlabel_changes = [\n ((0, 1), (1, 0), (2, 0)), # Set to black hair\n ((0, 0), (1, 1), (2, 0)), # Set to blonde hair\n ((0, 0), (1, 0), (2, 1)), # Set to brown hair\n ((3, -1),), # Flip gender\n ((4, -1),), # Age flip\n]\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample of domain translations\"\"\"\n val_imgs, val_labels = next(iter(val_dataloader))\n val_imgs = Variable(val_imgs.type(Tensor))\n val_labels = Variable(val_labels.type(Tensor))\n img_samples = None\n for i in range(10):\n img, label = val_imgs[i], val_labels[i]\n # Repeat for number of label changes\n imgs = img.repeat(c_dim, 1, 1, 1)\n labels = label.repeat(c_dim, 1)\n # Make changes to labels\n for sample_i, changes in enumerate(label_changes):\n for col, val in changes:\n labels[sample_i, col] = 1 - labels[sample_i, col] if val == -1 else val\n\n # Generate translations\n gen_imgs = generator(imgs, labels)\n # Concatenate images by width\n gen_imgs = torch.cat([x for x in gen_imgs.data], -1)\n img_sample = torch.cat((img.data, gen_imgs), -1)\n # Add as row to generated samples\n img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)\n\n save_image(img_samples.view(1, *img_samples.shape), \"images/%s.png\" % batches_done, normalize=True)\n\n\n# ----------\n# Training\n# ----------\n\nsaved_samples = []\nstart_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, (imgs, labels) in enumerate(dataloader):\n\n # Model inputs\n imgs = Variable(imgs.type(Tensor))\n labels = Variable(labels.type(Tensor))\n\n # Sample labels as generator inputs\n sampled_c = Variable(Tensor(np.random.randint(0, 2, (imgs.size(0), c_dim))))\n # Generate fake batch of images\n fake_imgs = generator(imgs, sampled_c)\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Real images\n real_validity, pred_cls = discriminator(imgs)\n # Fake images\n fake_validity, _ = discriminator(fake_imgs.detach())\n # Gradient penalty\n gradient_penalty = compute_gradient_penalty(discriminator, imgs.data, fake_imgs.data)\n # Adversarial loss\n loss_D_adv = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty\n # Classification loss\n loss_D_cls = criterion_cls(pred_cls, labels)\n # Total loss\n loss_D = loss_D_adv + lambda_cls * loss_D_cls\n\n loss_D.backward()\n optimizer_D.step()\n\n optimizer_G.zero_grad()\n\n # Every n_critic times update generator\n if i % opt.n_critic == 0:\n\n # -----------------\n # Train Generator\n # -----------------\n\n # Translate and reconstruct image\n gen_imgs = generator(imgs, sampled_c)\n recov_imgs = generator(gen_imgs, labels)\n # Discriminator evaluates translated image\n fake_validity, pred_cls = discriminator(gen_imgs)\n # Adversarial loss\n loss_G_adv = -torch.mean(fake_validity)\n # Classification loss\n loss_G_cls = criterion_cls(pred_cls, sampled_c)\n # Reconstruction loss\n loss_G_rec = criterion_cycle(recov_imgs, imgs)\n # Total loss\n loss_G = loss_G_adv + lambda_cls * loss_G_cls + lambda_rec * loss_G_rec\n\n loss_G.backward()\n optimizer_G.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - start_time) / (batches_done + 1))\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D adv: %f, aux: %f] [G loss: %f, adv: %f, aux: %f, cycle: %f] ETA: %s\"\n % (\n epoch,\n opt.n_epochs,\n i,\n len(dataloader),\n loss_D_adv.item(),\n loss_D_cls.item(),\n loss_G.item(),\n loss_G_adv.item(),\n loss_G_cls.item(),\n loss_G_rec.item(),\n time_left,\n )\n )\n\n # If at sample interval sample and save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(generator.state_dict(), \"saved_models/generator_%d.pth\" % epoch)\n torch.save(discriminator.state_dict(), \"saved_models/discriminator_%d.pth\" % epoch)\n"
},
{
"path": "implementations/unit/datasets.py",
"content": "import glob\nimport random\nimport os\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n def __init__(self, root, transforms_=None, unaligned=False, mode=\"train\"):\n self.transform = transforms.Compose(transforms_)\n self.unaligned = unaligned\n\n self.files_A = sorted(glob.glob(os.path.join(root, \"%s/A\" % mode) + \"/*.*\"))\n self.files_B = sorted(glob.glob(os.path.join(root, \"%s/B\" % mode) + \"/*.*\"))\n\n def __getitem__(self, index):\n item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]))\n\n if self.unaligned:\n item_B = self.transform(Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)]))\n else:\n item_B = self.transform(Image.open(self.files_B[index % len(self.files_B)]))\n\n return {\"A\": item_A, \"B\": item_B}\n\n def __len__(self):\n return max(len(self.files_A), len(self.files_B))\n"
},
{
"path": "implementations/unit/models.py",
"content": "import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torch.autograd import Variable\nimport numpy as np\n\n\ndef weights_init_normal(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass LambdaLR:\n def __init__(self, n_epochs, offset, decay_start_epoch):\n assert (n_epochs - decay_start_epoch) > 0, \"Decay must start before the training session ends!\"\n self.n_epochs = n_epochs\n self.offset = offset\n self.decay_start_epoch = decay_start_epoch\n\n def step(self, epoch):\n return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)\n\n\n##############################\n# RESNET\n##############################\n\n\nclass ResidualBlock(nn.Module):\n def __init__(self, features):\n super(ResidualBlock, self).__init__()\n\n conv_block = [\n nn.ReflectionPad2d(1),\n nn.Conv2d(features, features, 3),\n nn.InstanceNorm2d(features),\n nn.ReLU(inplace=True),\n nn.ReflectionPad2d(1),\n nn.Conv2d(features, features, 3),\n nn.InstanceNorm2d(features),\n ]\n\n self.conv_block = nn.Sequential(*conv_block)\n\n def forward(self, x):\n return x + self.conv_block(x)\n\n\nclass Encoder(nn.Module):\n def __init__(self, in_channels=3, dim=64, n_downsample=2, shared_block=None):\n super(Encoder, self).__init__()\n\n # Initial convolution block\n layers = [\n nn.ReflectionPad2d(3),\n nn.Conv2d(in_channels, dim, 7),\n nn.InstanceNorm2d(64),\n nn.LeakyReLU(0.2, inplace=True),\n ]\n\n # Downsampling\n for _ in range(n_downsample):\n layers += [\n nn.Conv2d(dim, dim * 2, 4, stride=2, padding=1),\n nn.InstanceNorm2d(dim * 2),\n nn.ReLU(inplace=True),\n ]\n dim *= 2\n\n # Residual blocks\n for _ in range(3):\n layers += [ResidualBlock(dim)]\n\n self.model_blocks = nn.Sequential(*layers)\n self.shared_block = shared_block\n\n def reparameterization(self, mu):\n Tensor = torch.cuda.FloatTensor if mu.is_cuda else torch.FloatTensor\n z = Variable(Tensor(np.random.normal(0, 1, mu.shape)))\n return z + mu\n\n def forward(self, x):\n x = self.model_blocks(x)\n mu = self.shared_block(x)\n z = self.reparameterization(mu)\n return mu, z\n\n\nclass Generator(nn.Module):\n def __init__(self, out_channels=3, dim=64, n_upsample=2, shared_block=None):\n super(Generator, self).__init__()\n\n self.shared_block = shared_block\n\n layers = []\n dim = dim * 2 ** n_upsample\n # Residual blocks\n for _ in range(3):\n layers += [ResidualBlock(dim)]\n\n # Upsampling\n for _ in range(n_upsample):\n layers += [\n nn.ConvTranspose2d(dim, dim // 2, 4, stride=2, padding=1),\n nn.InstanceNorm2d(dim // 2),\n nn.LeakyReLU(0.2, inplace=True),\n ]\n dim = dim // 2\n\n # Output layer\n layers += [nn.ReflectionPad2d(3), nn.Conv2d(dim, out_channels, 7), nn.Tanh()]\n\n self.model_blocks = nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.shared_block(x)\n x = self.model_blocks(x)\n return x\n\n\n##############################\n# Discriminator\n##############################\n\n\nclass Discriminator(nn.Module):\n def __init__(self, input_shape):\n super(Discriminator, self).__init__()\n channels, height, width = input_shape\n # Calculate output of image discriminator (PatchGAN)\n self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)\n\n def discriminator_block(in_filters, out_filters, normalize=True):\n \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n if normalize:\n layers.append(nn.InstanceNorm2d(out_filters))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *discriminator_block(channels, 64, normalize=False),\n *discriminator_block(64, 128),\n *discriminator_block(128, 256),\n *discriminator_block(256, 512),\n nn.Conv2d(512, 1, 3, padding=1)\n )\n\n def forward(self, img):\n return self.model(img)\n"
},
{
"path": "implementations/unit/unit.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport datetime\nimport time\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nfrom models import *\nfrom datasets import *\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--dataset_name\", type=str, default=\"apple2orange\", help=\"name of the dataset\")\nparser.add_argument(\"--batch_size\", type=int, default=1, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0001, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--img_height\", type=int, default=256, help=\"size of image height\")\nparser.add_argument(\"--img_width\", type=int, default=256, help=\"size of image width\")\nparser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\nparser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving generator samples\")\nparser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between saving model checkpoints\")\nparser.add_argument(\"--n_downsample\", type=int, default=2, help=\"number downsampling layers in encoder\")\nparser.add_argument(\"--dim\", type=int, default=64, help=\"number of filters in first encoder layer\")\nopt = parser.parse_args()\nprint(opt)\n\ncuda = True if torch.cuda.is_available() else False\n\n# Create sample and checkpoint directories\nos.makedirs(\"images/%s\" % opt.dataset_name, exist_ok=True)\nos.makedirs(\"saved_models/%s\" % opt.dataset_name, exist_ok=True)\n\n# Losses\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_pixel = torch.nn.L1Loss()\n\ninput_shape = (opt.channels, opt.img_height, opt.img_width)\n\n# Dimensionality (channel-wise) of image embedding\nshared_dim = opt.dim * 2 ** opt.n_downsample\n\n# Initialize generator and discriminator\nshared_E = ResidualBlock(features=shared_dim)\nE1 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, shared_block=shared_E)\nE2 = Encoder(dim=opt.dim, n_downsample=opt.n_downsample, shared_block=shared_E)\nshared_G = ResidualBlock(features=shared_dim)\nG1 = Generator(dim=opt.dim, n_upsample=opt.n_downsample, shared_block=shared_G)\nG2 = Generator(dim=opt.dim, n_upsample=opt.n_downsample, shared_block=shared_G)\nD1 = Discriminator(input_shape)\nD2 = Discriminator(input_shape)\n\nif cuda:\n E1 = E1.cuda()\n E2 = E2.cuda()\n G1 = G1.cuda()\n G2 = G2.cuda()\n D1 = D1.cuda()\n D2 = D2.cuda()\n criterion_GAN.cuda()\n criterion_pixel.cuda()\n\nif opt.epoch != 0:\n # Load pretrained models\n E1.load_state_dict(torch.load(\"saved_models/%s/E1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n E2.load_state_dict(torch.load(\"saved_models/%s/E2_%d.pth\" % (opt.dataset_name, opt.epoch)))\n G1.load_state_dict(torch.load(\"saved_models/%s/G1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n G2.load_state_dict(torch.load(\"saved_models/%s/G2_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D1.load_state_dict(torch.load(\"saved_models/%s/D1_%d.pth\" % (opt.dataset_name, opt.epoch)))\n D2.load_state_dict(torch.load(\"saved_models/%s/D2_%d.pth\" % (opt.dataset_name, opt.epoch)))\nelse:\n # Initialize weights\n E1.apply(weights_init_normal)\n E2.apply(weights_init_normal)\n G1.apply(weights_init_normal)\n G2.apply(weights_init_normal)\n D1.apply(weights_init_normal)\n D2.apply(weights_init_normal)\n\n# Loss weights\nlambda_0 = 10 # GAN\nlambda_1 = 0.1 # KL (encoded images)\nlambda_2 = 100 # ID pixel-wise\nlambda_3 = 0.1 # KL (encoded translated images)\nlambda_4 = 100 # Cycle pixel-wise\n\n# Optimizers\noptimizer_G = torch.optim.Adam(\n itertools.chain(E1.parameters(), E2.parameters(), G1.parameters(), G2.parameters()),\n lr=opt.lr,\n betas=(opt.b1, opt.b2),\n)\noptimizer_D1 = torch.optim.Adam(D1.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D2 = torch.optim.Adam(D2.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\n# Learning rate update schedulers\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n optimizer_G, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D1 = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D1, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\nlr_scheduler_D2 = torch.optim.lr_scheduler.LambdaLR(\n optimizer_D2, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n)\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\n# Image transformations\ntransforms_ = [\n transforms.Resize(int(opt.img_height * 1.12), Image.BICUBIC),\n transforms.RandomCrop((opt.img_height, opt.img_width)),\n transforms.RandomHorizontalFlip(),\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n]\n\n# Training data loader\ndataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True),\n batch_size=opt.batch_size,\n shuffle=True,\n num_workers=opt.n_cpu,\n)\n# Test data loader\nval_dataloader = DataLoader(\n ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True, mode=\"test\"),\n batch_size=5,\n shuffle=True,\n num_workers=1,\n)\n\n\ndef sample_images(batches_done):\n \"\"\"Saves a generated sample from the test set\"\"\"\n imgs = next(iter(val_dataloader))\n X1 = Variable(imgs[\"A\"].type(Tensor))\n X2 = Variable(imgs[\"B\"].type(Tensor))\n _, Z1 = E1(X1)\n _, Z2 = E2(X2)\n fake_X1 = G1(Z2)\n fake_X2 = G2(Z1)\n img_sample = torch.cat((X1.data, fake_X2.data, X2.data, fake_X1.data), 0)\n save_image(img_sample, \"images/%s/%s.png\" % (opt.dataset_name, batches_done), nrow=5, normalize=True)\n\n\ndef compute_kl(mu):\n mu_2 = torch.pow(mu, 2)\n loss = torch.mean(mu_2)\n return loss\n\n\n# ----------\n# Training\n# ----------\n\nprev_time = time.time()\nfor epoch in range(opt.epoch, opt.n_epochs):\n for i, batch in enumerate(dataloader):\n\n # Set model input\n X1 = Variable(batch[\"A\"].type(Tensor))\n X2 = Variable(batch[\"B\"].type(Tensor))\n\n # Adversarial ground truths\n valid = Variable(Tensor(np.ones((X1.size(0), *D1.output_shape))), requires_grad=False)\n fake = Variable(Tensor(np.zeros((X1.size(0), *D1.output_shape))), requires_grad=False)\n\n # -------------------------------\n # Train Encoders and Generators\n # -------------------------------\n\n optimizer_G.zero_grad()\n\n # Get shared latent representation\n mu1, Z1 = E1(X1)\n mu2, Z2 = E2(X2)\n\n # Reconstruct images\n recon_X1 = G1(Z1)\n recon_X2 = G2(Z2)\n\n # Translate images\n fake_X1 = G1(Z2)\n fake_X2 = G2(Z1)\n\n # Cycle translation\n mu1_, Z1_ = E1(fake_X1)\n mu2_, Z2_ = E2(fake_X2)\n cycle_X1 = G1(Z2_)\n cycle_X2 = G2(Z1_)\n\n # Losses\n loss_GAN_1 = lambda_0 * criterion_GAN(D1(fake_X1), valid)\n loss_GAN_2 = lambda_0 * criterion_GAN(D2(fake_X2), valid)\n loss_KL_1 = lambda_1 * compute_kl(mu1)\n loss_KL_2 = lambda_1 * compute_kl(mu2)\n loss_ID_1 = lambda_2 * criterion_pixel(recon_X1, X1)\n loss_ID_2 = lambda_2 * criterion_pixel(recon_X2, X2)\n loss_KL_1_ = lambda_3 * compute_kl(mu1_)\n loss_KL_2_ = lambda_3 * compute_kl(mu2_)\n loss_cyc_1 = lambda_4 * criterion_pixel(cycle_X1, X1)\n loss_cyc_2 = lambda_4 * criterion_pixel(cycle_X2, X2)\n\n # Total loss\n loss_G = (\n loss_KL_1\n + loss_KL_2\n + loss_ID_1\n + loss_ID_2\n + loss_GAN_1\n + loss_GAN_2\n + loss_KL_1_\n + loss_KL_2_\n + loss_cyc_1\n + loss_cyc_2\n )\n\n loss_G.backward()\n optimizer_G.step()\n\n # -----------------------\n # Train Discriminator 1\n # -----------------------\n\n optimizer_D1.zero_grad()\n\n loss_D1 = criterion_GAN(D1(X1), valid) + criterion_GAN(D1(fake_X1.detach()), fake)\n\n loss_D1.backward()\n optimizer_D1.step()\n\n # -----------------------\n # Train Discriminator 2\n # -----------------------\n\n optimizer_D2.zero_grad()\n\n loss_D2 = criterion_GAN(D2(X2), valid) + criterion_GAN(D2(fake_X2.detach()), fake)\n\n loss_D2.backward()\n optimizer_D2.step()\n\n # --------------\n # Log Progress\n # --------------\n\n # Determine approximate time left\n batches_done = epoch * len(dataloader) + i\n batches_left = opt.n_epochs * len(dataloader) - batches_done\n time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n prev_time = time.time()\n\n # Print log\n sys.stdout.write(\n \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] ETA: %s\"\n % (epoch, opt.n_epochs, i, len(dataloader), (loss_D1 + loss_D2).item(), loss_G.item(), time_left)\n )\n\n # If at sample interval save image\n if batches_done % opt.sample_interval == 0:\n sample_images(batches_done)\n\n # Update learning rates\n lr_scheduler_G.step()\n lr_scheduler_D1.step()\n lr_scheduler_D2.step()\n\n if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n # Save model checkpoints\n torch.save(E1.state_dict(), \"saved_models/%s/E1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(E2.state_dict(), \"saved_models/%s/E2_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(G1.state_dict(), \"saved_models/%s/G1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(G2.state_dict(), \"saved_models/%s/G2_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D1.state_dict(), \"saved_models/%s/D1_%d.pth\" % (opt.dataset_name, epoch))\n torch.save(D2.state_dict(), \"saved_models/%s/D2_%d.pth\" % (opt.dataset_name, epoch))\n"
},
{
"path": "implementations/wgan/wgan.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.00005, help=\"learning rate\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--n_critic\", type=int, default=5, help=\"number of training steps for discriminator per iter\")\nparser.add_argument(\"--clip_value\", type=float, default=0.01, help=\"lower and upper clip value for disc. weights\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n )\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n validity = self.model(img_flat)\n return validity\n\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.RMSprop(generator.parameters(), lr=opt.lr)\noptimizer_D = torch.optim.RMSprop(discriminator.parameters(), lr=opt.lr)\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nbatches_done = 0\nfor epoch in range(opt.n_epochs):\n\n for i, (imgs, _) in enumerate(dataloader):\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n fake_imgs = generator(z).detach()\n # Adversarial loss\n loss_D = -torch.mean(discriminator(real_imgs)) + torch.mean(discriminator(fake_imgs))\n\n loss_D.backward()\n optimizer_D.step()\n\n # Clip weights of discriminator\n for p in discriminator.parameters():\n p.data.clamp_(-opt.clip_value, opt.clip_value)\n\n # Train the generator every n_critic iterations\n if i % opt.n_critic == 0:\n\n # -----------------\n # Train Generator\n # -----------------\n\n optimizer_G.zero_grad()\n\n # Generate a batch of images\n gen_imgs = generator(z)\n # Adversarial loss\n loss_G = -torch.mean(discriminator(gen_imgs))\n\n loss_G.backward()\n optimizer_G.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, batches_done % len(dataloader), len(dataloader), loss_D.item(), loss_G.item())\n )\n\n if batches_done % opt.sample_interval == 0:\n save_image(gen_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n batches_done += 1\n"
},
{
"path": "implementations/wgan_div/wgan_div.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.autograd as autograd\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--n_critic\", type=int, default=5, help=\"number of training steps for discriminator per iter\")\nparser.add_argument(\"--clip_value\", type=float, default=0.01, help=\"lower and upper clip value for disc. weights\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n )\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n validity = self.model(img_flat)\n return validity\n\n\nk = 2\np = 6\n\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n# ----------\n# Training\n# ----------\n\nbatches_done = 0\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor), requires_grad=True)\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n fake_imgs = generator(z)\n\n # Real images\n real_validity = discriminator(real_imgs)\n # Fake images\n fake_validity = discriminator(fake_imgs)\n\n # Compute W-div gradient penalty\n real_grad_out = Variable(Tensor(real_imgs.size(0), 1).fill_(1.0), requires_grad=False)\n real_grad = autograd.grad(\n real_validity, real_imgs, real_grad_out, create_graph=True, retain_graph=True, only_inputs=True\n )[0]\n real_grad_norm = real_grad.view(real_grad.size(0), -1).pow(2).sum(1) ** (p / 2)\n\n fake_grad_out = Variable(Tensor(fake_imgs.size(0), 1).fill_(1.0), requires_grad=False)\n fake_grad = autograd.grad(\n fake_validity, fake_imgs, fake_grad_out, create_graph=True, retain_graph=True, only_inputs=True\n )[0]\n fake_grad_norm = fake_grad.view(fake_grad.size(0), -1).pow(2).sum(1) ** (p / 2)\n\n div_gp = torch.mean(real_grad_norm + fake_grad_norm) * k / 2\n\n # Adversarial loss\n d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + div_gp\n\n d_loss.backward()\n optimizer_D.step()\n\n optimizer_G.zero_grad()\n\n # Train the generator every n_critic steps\n if i % opt.n_critic == 0:\n\n # -----------------\n # Train Generator\n # -----------------\n\n # Generate a batch of images\n fake_imgs = generator(z)\n # Loss measures generator's ability to fool the discriminator\n # Train on fake images\n fake_validity = discriminator(fake_imgs)\n g_loss = -torch.mean(fake_validity)\n\n g_loss.backward()\n optimizer_G.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n if batches_done % opt.sample_interval == 0:\n save_image(fake_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n\n batches_done += opt.n_critic\n"
},
{
"path": "implementations/wgan_gp/wgan_gp.py",
"content": "import argparse\nimport os\nimport numpy as np\nimport math\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.autograd as autograd\nimport torch\n\nos.makedirs(\"images\", exist_ok=True)\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\nparser.add_argument(\"--batch_size\", type=int, default=64, help=\"size of the batches\")\nparser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\nparser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\nparser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\nparser.add_argument(\"--latent_dim\", type=int, default=100, help=\"dimensionality of the latent space\")\nparser.add_argument(\"--img_size\", type=int, default=28, help=\"size of each image dimension\")\nparser.add_argument(\"--channels\", type=int, default=1, help=\"number of image channels\")\nparser.add_argument(\"--n_critic\", type=int, default=5, help=\"number of training steps for discriminator per iter\")\nparser.add_argument(\"--clip_value\", type=float, default=0.01, help=\"lower and upper clip value for disc. weights\")\nparser.add_argument(\"--sample_interval\", type=int, default=400, help=\"interval betwen image samples\")\nopt = parser.parse_args()\nprint(opt)\n\nimg_shape = (opt.channels, opt.img_size, opt.img_size)\n\ncuda = True if torch.cuda.is_available() else False\n\n\nclass Generator(nn.Module):\n def __init__(self):\n super(Generator, self).__init__()\n\n def block(in_feat, out_feat, normalize=True):\n layers = [nn.Linear(in_feat, out_feat)]\n if normalize:\n layers.append(nn.BatchNorm1d(out_feat, 0.8))\n layers.append(nn.LeakyReLU(0.2, inplace=True))\n return layers\n\n self.model = nn.Sequential(\n *block(opt.latent_dim, 128, normalize=False),\n *block(128, 256),\n *block(256, 512),\n *block(512, 1024),\n nn.Linear(1024, int(np.prod(img_shape))),\n nn.Tanh()\n )\n\n def forward(self, z):\n img = self.model(z)\n img = img.view(img.shape[0], *img_shape)\n return img\n\n\nclass Discriminator(nn.Module):\n def __init__(self):\n super(Discriminator, self).__init__()\n\n self.model = nn.Sequential(\n nn.Linear(int(np.prod(img_shape)), 512),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(512, 256),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Linear(256, 1),\n )\n\n def forward(self, img):\n img_flat = img.view(img.shape[0], -1)\n validity = self.model(img_flat)\n return validity\n\n\n# Loss weight for gradient penalty\nlambda_gp = 10\n\n# Initialize generator and discriminator\ngenerator = Generator()\ndiscriminator = Discriminator()\n\nif cuda:\n generator.cuda()\n discriminator.cuda()\n\n# Configure data loader\nos.makedirs(\"../../data/mnist\", exist_ok=True)\ndataloader = torch.utils.data.DataLoader(\n datasets.MNIST(\n \"../../data/mnist\",\n train=True,\n download=True,\n transform=transforms.Compose(\n [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n ),\n ),\n batch_size=opt.batch_size,\n shuffle=True,\n)\n\n# Optimizers\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n\nTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n\n\ndef compute_gradient_penalty(D, real_samples, fake_samples):\n \"\"\"Calculates the gradient penalty loss for WGAN GP\"\"\"\n # Random weight term for interpolation between real and fake samples\n alpha = Tensor(np.random.random((real_samples.size(0), 1, 1, 1)))\n # Get random interpolation between real and fake samples\n interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)\n d_interpolates = D(interpolates)\n fake = Variable(Tensor(real_samples.shape[0], 1).fill_(1.0), requires_grad=False)\n # Get gradient w.r.t. interpolates\n gradients = autograd.grad(\n outputs=d_interpolates,\n inputs=interpolates,\n grad_outputs=fake,\n create_graph=True,\n retain_graph=True,\n only_inputs=True,\n )[0]\n gradients = gradients.view(gradients.size(0), -1)\n gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n return gradient_penalty\n\n\n# ----------\n# Training\n# ----------\n\nbatches_done = 0\nfor epoch in range(opt.n_epochs):\n for i, (imgs, _) in enumerate(dataloader):\n\n # Configure input\n real_imgs = Variable(imgs.type(Tensor))\n\n # ---------------------\n # Train Discriminator\n # ---------------------\n\n optimizer_D.zero_grad()\n\n # Sample noise as generator input\n z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n\n # Generate a batch of images\n fake_imgs = generator(z)\n\n # Real images\n real_validity = discriminator(real_imgs)\n # Fake images\n fake_validity = discriminator(fake_imgs)\n # Gradient penalty\n gradient_penalty = compute_gradient_penalty(discriminator, real_imgs.data, fake_imgs.data)\n # Adversarial loss\n d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty\n\n d_loss.backward()\n optimizer_D.step()\n\n optimizer_G.zero_grad()\n\n # Train the generator every n_critic steps\n if i % opt.n_critic == 0:\n\n # -----------------\n # Train Generator\n # -----------------\n\n # Generate a batch of images\n fake_imgs = generator(z)\n # Loss measures generator's ability to fool the discriminator\n # Train on fake images\n fake_validity = discriminator(fake_imgs)\n g_loss = -torch.mean(fake_validity)\n\n g_loss.backward()\n optimizer_G.step()\n\n print(\n \"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\"\n % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())\n )\n\n if batches_done % opt.sample_interval == 0:\n save_image(fake_imgs.data[:25], \"images/%d.png\" % batches_done, nrow=5, normalize=True)\n\n batches_done += opt.n_critic\n"
},
{
"path": "requirements.txt",
"content": "torch>=0.4.0\ntorchvision\nmatplotlib\nnumpy\nscipy\npillow\nurllib3\nscikit-image\n"
}
]![](\"assets/logo.png\"){kind=logo}
![](\"assets/acgan.gif\")
\n![](\"assets/bicyclegan_architecture.jpg\")
\n![](\"assets/bicyclegan.png\")
\n![](\"assets/cluster_gan.gif\")
\n![](\"assets/cgan.gif\")
\n