""" --- title: Cycle GAN summary: > A simple PyTorch implementation/tutorial of Cycle GAN introduced in paper Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks. --- # Cycle GAN This is a [PyTorch](https://pytorch.org) implementation/tutorial of paper [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593). I've taken pieces of code from [eriklindernoren/PyTorch-GAN](https://github.com/eriklindernoren/PyTorch-GAN). It is a very good resource if you want to checkout other GAN variations too. Cycle GAN does image-to-image translation. It trains a model to translate an image from given distribution to another, say images of class A and B Images of a certain distribution could be things like images of a certain style, or nature. The models do not need paired images between A and B. Just a set of images of each class is enough. This works very well on changing between image styles, lighting changes, pattern changes, etc. For example, changing summer to winter, painting style to photos, and horses to zebras. Cycle GAN trains two generator models and two discriminator models. One generator translates images from A to B and the other from B to A. The discriminators test whether the generated images look real. This file contains the model code as well as training code. We also have a Google Colab notebook. [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lab-ml/nn/blob/master/labml_nn/gan/cycle_gan.ipynb) [![View Run](https://img.shields.io/badge/labml-experiment-brightgreen)](https://web.lab-ml.com/run?uuid=93b11a665d6811ebaac80242ac1c0002) """ import itertools import random import zipfile from typing import Tuple import torch import torch.nn as nn import torchvision.transforms as transforms from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision.utils import make_grid from labml import lab, tracker, experiment, monit from labml.configs import BaseConfigs from labml.utils.download import download_file from labml.utils.pytorch import get_modules from labml_helpers.device import DeviceConfigs from labml_helpers.module import Module class GeneratorResNet(Module): """ The generator is a residual network. """ def __init__(self, input_channels: int, n_residual_blocks: int): super().__init__() # This first block runs a $7\times7$ convolution and maps the image to # a feature map. # The output feature map has same height and width because we have # a padding of $3$. # Reflection padding is used because it gives better image quality at edges. # # `inplace=True` in `ReLU` saves a little bit of memory. out_features = 64 layers = [ nn.Conv2d(input_channels, out_features, kernel_size=7, padding=3, padding_mode='reflect'), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True), ] in_features = out_features # We down-sample with two $3 \times 3$ convolutions # with stride of 2 for _ in range(2): out_features *= 2 layers += [ nn.Conv2d(in_features, out_features, kernel_size=3, stride=2, padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True), ] in_features = out_features # We take this through `n_residual_blocks`. # This module is defined below. for _ in range(n_residual_blocks): layers += [ResidualBlock(out_features)] # Then the resulting feature map is up-sampled # to match the original image height and width. for _ in range(2): out_features //= 2 layers += [ nn.Upsample(scale_factor=2), nn.Conv2d(in_features, out_features, kernel_size=3, stride=1, padding=1), nn.InstanceNorm2d(out_features), nn.ReLU(inplace=True), ] in_features = out_features # Finally we map the feature map to an RGB image layers += [nn.Conv2d(out_features, input_channels, 7, padding=3, padding_mode='reflect'), nn.Tanh()] # Create a sequential module with the layers self.layers = nn.Sequential(*layers) # Initialize weights to $\mathcal{N}(0, 0.2)$ self.apply(weights_init_normal) def __call__(self, x): return self.layers(x) class ResidualBlock(Module): """ This is the residual block, with two convolution layers. """ def __init__(self, in_features: int): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_features, in_features, kernel_size=3, padding=1, padding_mode='reflect'), nn.InstanceNorm2d(in_features), nn.ReLU(inplace=True), nn.Conv2d(in_features, in_features, kernel_size=3, padding=1, padding_mode='reflect'), nn.InstanceNorm2d(in_features), nn.ReLU(inplace=True), ) def __call__(self, x: torch.Tensor): return x + self.block(x) class Discriminator(Module): """ This is the discriminator. """ def __init__(self, input_shape: Tuple[int, int, int]): super().__init__() channels, height, width = input_shape # Output of the discriminator is also map of probabilities* # whether each region of the image is real or generated self.output_shape = (1, height // 2 ** 4, width // 2 ** 4) self.layers = nn.Sequential( # Each of these blocks will shrink the height and width by a factor of 2 DiscriminatorBlock(channels, 64, normalize=False), DiscriminatorBlock(64, 128), DiscriminatorBlock(128, 256), DiscriminatorBlock(256, 512), # Zero pad on top and left to keep the output height and width same # with the $4 \times 4$ kernel nn.ZeroPad2d((1, 0, 1, 0)), nn.Conv2d(512, 1, kernel_size=4, padding=1) ) # Initialize weights to $\mathcal{N}(0, 0.2)$ self.apply(weights_init_normal) def forward(self, img): return self.layers(img) class DiscriminatorBlock(Module): """ This is the discriminator block module. It does a convolution, an optional normalization, and a leaky relu. It shrinks the height and width of the input feature map by half. """ def __init__(self, in_filters: int, out_filters: int, normalize: bool = True): super().__init__() layers = [nn.Conv2d(in_filters, out_filters, kernel_size=4, stride=2, padding=1)] if normalize: layers.append(nn.InstanceNorm2d(out_filters)) layers.append(nn.LeakyReLU(0.2, inplace=True)) self.layers = nn.Sequential(*layers) def __call__(self, x: torch.Tensor): return self.layers(x) def weights_init_normal(m): """ Initialize convolution layer weights to $\mathcal{N}(0, 0.2)$ """ classname = m.__class__.__name__ if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) def load_image(path: str): """ Loads an image and change to RGB if in grey-scale. """ image = Image.open(path) if image.mode != 'RGB': image = Image.new("RGB", image.size).paste(image) return image class ImageDataset(Dataset): """ ### Dataset to load images """ @staticmethod def download(dataset_name: str): """ #### Download dataset and extract data """ # URL url = f'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/{dataset_name}.zip' # Download folder root = lab.get_data_path() / 'cycle_gan' if not root.exists(): root.mkdir(parents=True) # Download destination archive = root / f'{dataset_name}.zip' # Download file (generally ~100MB) download_file(url, archive) # Extract the archive with zipfile.ZipFile(archive, 'r') as f: f.extractall(root) def __init__(self, dataset_name: str, transforms_, mode: str): """ #### Initialize the dataset * `dataset_name` is the name of the dataset * `transforms_` is the set of image transforms * `mode` is either `train` or `test` """ # Dataset path root = lab.get_data_path() / 'cycle_gan' / dataset_name # Download if missing if not root.exists(): self.download(dataset_name) # Image transforms self.transform = transforms.Compose(transforms_) # Get image paths path_a = root / f'{mode}A' path_b = root / f'{mode}B' self.files_a = sorted(str(f) for f in path_a.iterdir()) self.files_b = sorted(str(f) for f in path_b.iterdir()) def __getitem__(self, index): # Return a pair of images. # These pairs get batched together, and they do not act like pairs in training # So it is kind of ok that we always keep giving the same pair. return {"x": self.transform(load_image(self.files_a[index % len(self.files_a)])), "y": self.transform(load_image(self.files_b[index % len(self.files_b)]))} def __len__(self): # Number of images in the dataset return max(len(self.files_a), len(self.files_b)) class ReplayBuffer: """ ### Replay Buffer Replay buffer is used to train the discriminator. Generated images are added to the replay buffer and sampled from it. The replay buffer returns the newly added image with a probability of $0.5$. Otherwise it sends an older generated image and and replaces the older image with the new generated image. This is done to reduce model oscillation. """ def __init__(self, max_size: int = 50): self.max_size = max_size self.data = [] def push_and_pop(self, data: torch.Tensor): """Add/retrieve an image""" data = data.detach() res = [] for element in data: if len(self.data) < self.max_size: self.data.append(element) res.append(element) else: if random.uniform(0, 1) > 0.5: i = random.randint(0, self.max_size - 1) res.append(self.data[i].clone()) self.data[i] = element else: res.append(element) return torch.stack(res) class Configs(BaseConfigs): """## Configurations""" # `DeviceConfigs` will pick a GPU if available device: torch.device = DeviceConfigs() # Hyper-parameters epochs: int = 200 dataset_name: str = 'monet2photo' batch_size: int = 1 data_loader_workers = 8 learning_rate = 0.0002 adam_betas = (0.5, 0.999) decay_start = 100 # The paper suggests using a least-squares loss instead of # negative log-likelihood, at it is found to be more stable. gan_loss = torch.nn.MSELoss() # L1 loss is used for cycle loss and identity loss cycle_loss = torch.nn.L1Loss() identity_loss = torch.nn.L1Loss() # Image dimensions img_height = 256 img_width = 256 img_channels = 3 # Number of residual blocks in the generator n_residual_blocks = 9 # Loss coefficients cyclic_loss_coefficient = 10.0 identity_loss_coefficient = 5. sample_interval = 500 # Models generator_xy: GeneratorResNet generator_yx: GeneratorResNet discriminator_x: Discriminator discriminator_y: Discriminator # Optimizers generator_optimizer: torch.optim.Adam discriminator_optimizer: torch.optim.Adam # Learning rate schedules generator_lr_scheduler: torch.optim.lr_scheduler.LambdaLR discriminator_lr_scheduler: torch.optim.lr_scheduler.LambdaLR # Data loaders dataloader: DataLoader valid_dataloader: DataLoader def sample_images(self, n: int): """Generate samples from test set and save them""" batch = next(iter(self.valid_dataloader)) self.generator_xy.eval() self.generator_yx.eval() with torch.no_grad(): data_x, data_y = batch['x'].to(self.generator_xy.device), batch['y'].to(self.generator_yx.device) gen_y = self.generator_xy(data_x) gen_x = self.generator_yx(data_y) # Arrange images along x-axis data_x = make_grid(data_x, nrow=5, normalize=True) data_y = make_grid(data_y, nrow=5, normalize=True) gen_x = make_grid(gen_x, nrow=5, normalize=True) gen_y = make_grid(gen_y, nrow=5, normalize=True) # Arrange images along y-axis image_grid = torch.cat((data_x, gen_y, data_y, gen_x), 1) # Show samples plot_image(image_grid) def initialize(self): """ ## Initialize models and data loaders """ input_shape = (self.img_channels, self.img_height, self.img_width) # Create the models self.generator_xy = GeneratorResNet(self.img_channels, self.n_residual_blocks).to(self.device) self.generator_yx = GeneratorResNet(self.img_channels, self.n_residual_blocks).to(self.device) self.discriminator_x = Discriminator(input_shape).to(self.device) self.discriminator_y = Discriminator(input_shape).to(self.device) # Create the optmizers self.generator_optimizer = torch.optim.Adam( itertools.chain(self.generator_xy.parameters(), self.generator_yx.parameters()), lr=self.learning_rate, betas=self.adam_betas) self.discriminator_optimizer = torch.optim.Adam( itertools.chain(self.discriminator_x.parameters(), self.discriminator_y.parameters()), lr=self.learning_rate, betas=self.adam_betas) # Create the learning rate schedules. # The learning rate stars flat until `decay_start` epochs, # and then linearly reduces to $0$ at end of training. decay_epochs = self.epochs - self.decay_start self.generator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR( self.generator_optimizer, lr_lambda=lambda e: 1.0 - max(0, e - self.decay_start) / decay_epochs) self.discriminator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR( self.discriminator_optimizer, lr_lambda=lambda e: 1.0 - max(0, e - self.decay_start) / decay_epochs) # Image transformations transforms_ = [ transforms.Resize(int(self.img_height * 1.12), Image.BICUBIC), transforms.RandomCrop((self.img_height, self.img_width)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] # Training data loader self.dataloader = DataLoader( ImageDataset(self.dataset_name, transforms_, 'train'), batch_size=self.batch_size, shuffle=True, num_workers=self.data_loader_workers, ) # Validation data loader self.valid_dataloader = DataLoader( ImageDataset(self.dataset_name, transforms_, "test"), batch_size=5, shuffle=True, num_workers=self.data_loader_workers, ) def run(self): """ ## Training We aim to solve: $$G^{*}, F^{*} = \arg \min_{G,F} \max_{D_X, D_Y} \mathcal{L}(G, F, D_X, D_Y)$$ where, $G$ translates images from $X \rightarrow Y$, $F$ translates images from $Y \rightarrow X$, $D_X$ tests if images are from $X$ space, $D_Y$ tests if images are from $Y$ space, and \begin{align} \mathcal{L}(G, F, D_X, D_Y) &= \mathcal{L}_{GAN}(G, D_Y, X, Y) \\ &+ \mathcal{L}_{GAN}(F, D_X, Y, X) \\ &+ \lambda_1 \mathcal{L}_{cyc}(G, F) \\ &+ \lambda_2 \mathcal{L}_{identity}(G, F) \\ \\ \mathcal{L}_{GAN}(G, F, D_Y, X, Y) &= \mathbb{E}_{y \sim p_{data}(y)} \Big[log D_Y(y)\Big] \\ &+ \mathbb{E}_{x \sim p_{data}(x)} \bigg[log\Big(1 - D_Y(G(x))\Big)\bigg] \\ &+ \mathbb{E}_{x \sim p_{data}(x)} \Big[log D_X(x)\Big] \\ &+ \mathbb{E}_{y \sim p_{data}(y)} \bigg[log\Big(1 - D_X(F(y))\Big)\bigg] \\ \\ \mathcal{L}_{cyc}(G, F) &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(G(x)) - x \lVert_1\Big] \\ &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(F(y)) - y \rVert_1\Big] \\ \\ \mathcal{L}_{identity}(G, F) &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(x) - x \lVert_1\Big] \\ &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(y) - y \rVert_1\Big] \\ \end{align} $\mathcal{L}_{GAN}$ is the generative adversarial loss from the original GAN paper. $\mathcal{L}_{cyc}$ is the cyclic loss, where we try to get $F(G(x))$ to be similar to $x$, and $G(F(y))$ to be similar to $y$. Basically if the two generators (transformations) are applied in series it should give back the original image. This is the main contribution of this paper. It train the generators to generate an image of the other distribution that is similar to the original image. Without this loss $G(x)$ could generate anything that's from the distribution of $Y$. Now it needs to generate something from the distribution of $Y$ but still have properties of $x$, so that $F(G(x)$ can re-generate something like $x$. $\mathcal{L}_{cyc}$ is the identity loss. This was used to encourage the mapping to preserve color composition between the input and the output. To solve $G^{\*}, F^{\*}$, discriminators $D_X$ and $D_Y$ should **ascend** on the gradient, \begin{align} \nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m &\Bigg[ \log D_Y\Big(y^{(i)}\Big) \\ &+ \log \Big(1 - D_Y\Big(G\Big(x^{(i)}\Big)\Big)\Big) \\ &+ \log D_X\Big(x^{(i)}\Big) \\ & +\log\Big(1 - D_X\Big(F\Big(y^{(i)}\Big)\Big)\Big) \Bigg] \end{align} That is descend on *negative* log-likelihood loss. In order to stabilize the training the negative log- likelihood objective was replaced by a least-squared loss - the least-squared error of discriminator labelling real images with 1, and generated images with 0. So we want to descend on the gradient, \begin{align} \nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m &\Bigg[ \bigg(D_Y\Big(y^{(i)}\Big) - 1\bigg)^2 \\ &+ D_Y\Big(G\Big(x^{(i)}\Big)\Big)^2 \\ &+ \bigg(D_X\Big(x^{(i)}\Big) - 1\bigg)^2 \\ &+ D_X\Big(F\Big(y^{(i)}\Big)\Big)^2 \Bigg] \end{align} We use least-squares for generators also. The generators should *descend* on the gradient, \begin{align} \nabla_{\theta_{F, G}} \frac{1}{m} \sum_{i=1}^m &\Bigg[ \bigg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1\bigg)^2 \\ &+ \bigg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1\bigg)^2 \\ &+ \mathcal{L}_{cyc}(G, F) + \mathcal{L}_{identity}(G, F) \Bigg] \end{align} We use `generator_xy` for $G$ and `generator_yx$ for $F$. We use `discriminator_x$ for $D_X$ and `discriminator_y` for $D_Y$. """ # Replay buffers to keep generated samples gen_x_buffer = ReplayBuffer() gen_y_buffer = ReplayBuffer() # Loop through epochs for epoch in monit.loop(self.epochs): # Loop through the dataset for i, batch in monit.enum('Train', self.dataloader): # Move images to the device data_x, data_y = batch['x'].to(self.device), batch['y'].to(self.device) # true labels equal to $1$ true_labels = torch.ones(data_x.size(0), *self.discriminator_x.output_shape, device=self.device, requires_grad=False) # false labels equal to $0$ false_labels = torch.zeros(data_x.size(0), *self.discriminator_x.output_shape, device=self.device, requires_grad=False) # Train the generators. # This returns the generated images. gen_x, gen_y = self.optimize_generators(data_x, data_y, true_labels) # Train discriminators self.optimize_discriminator(data_x, data_y, gen_x_buffer.push_and_pop(gen_x), gen_y_buffer.push_and_pop(gen_y), true_labels, false_labels) # Save training statistics and increment the global step counter tracker.save() tracker.add_global_step(max(len(data_x), len(data_y))) # Save images at intervals batches_done = epoch * len(self.dataloader) + i if batches_done % self.sample_interval == 0: # Save models when sampling images experiment.save_checkpoint() # Sample images self.sample_images(batches_done) # Update learning rates self.generator_lr_scheduler.step() self.discriminator_lr_scheduler.step() # New line tracker.new_line() def optimize_generators(self, data_x: torch.Tensor, data_y: torch.Tensor, true_labels: torch.Tensor): """ ### Optimize the generators with identity, gan and cycle losses. """ # Change to training mode self.generator_xy.train() self.generator_yx.train() # Identity loss # $$\lVert F(G(x^{(i)})) - x^{(i)} \lVert_1\ # \lVert G(F(y^{(i)})) - y^{(i)} \rVert_1$$ loss_identity = (self.identity_loss(self.generator_yx(data_x), data_x) + self.identity_loss(self.generator_xy(data_y), data_y)) # Generate images $G(x)$ and $F(y)$ gen_y = self.generator_xy(data_x) gen_x = self.generator_yx(data_y) # GAN loss # $$\bigg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1\bigg)^2 # + \bigg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1\bigg)^2$$ loss_gan = (self.gan_loss(self.discriminator_y(gen_y), true_labels) + self.gan_loss(self.discriminator_x(gen_x), true_labels)) # Cycle loss # $$ # \lVert F(G(x^{(i)})) - x^{(i)} \lVert_1 + # \lVert G(F(y^{(i)})) - y^{(i)} \rVert_1 # $$ loss_cycle = (self.cycle_loss(self.generator_yx(gen_y), data_x) + self.cycle_loss(self.generator_xy(gen_x), data_y)) # Total loss loss_generator = (loss_gan + self.cyclic_loss_coefficient * loss_cycle + self.identity_loss_coefficient * loss_identity) # Take a step in the optimizer self.generator_optimizer.zero_grad() loss_generator.backward() self.generator_optimizer.step() # Log losses tracker.add({'loss.generator': loss_generator, 'loss.generator.cycle': loss_cycle, 'loss.generator.gan': loss_gan, 'loss.generator.identity': loss_identity}) # Return generated images return gen_x, gen_y def optimize_discriminator(self, data_x: torch.Tensor, data_y: torch.Tensor, gen_x: torch.Tensor, gen_y: torch.Tensor, true_labels: torch.Tensor, false_labels: torch.Tensor): """ ### Optimize the discriminators with gan loss. """ # GAN Loss # \begin{align} # \bigg(D_Y\Big(y ^ {(i)}\Big) - 1\bigg) ^ 2 # + D_Y\Big(G\Big(x ^ {(i)}\Big)\Big) ^ 2 + \\ # \bigg(D_X\Big(x ^ {(i)}\Big) - 1\bigg) ^ 2 # + D_X\Big(F\Big(y ^ {(i)}\Big)\Big) ^ 2 # \end{align} loss_discriminator = (self.gan_loss(self.discriminator_x(data_x), true_labels) + self.gan_loss(self.discriminator_x(gen_x), false_labels) + self.gan_loss(self.discriminator_y(data_y), true_labels) + self.gan_loss(self.discriminator_y(gen_y), false_labels)) # Take a step in the optimizer self.discriminator_optimizer.zero_grad() loss_discriminator.backward() self.discriminator_optimizer.step() # Log losses tracker.add({'loss.discriminator': loss_discriminator}) def train(): """ ## Train Cycle GAN """ # Create configurations conf = Configs() # Create an experiment experiment.create(name='cycle_gan') # Calculate configurations. # It will calculate `conf.run` and all other configs required by it. experiment.configs(conf, {'dataset_name': 'summer2winter_yosemite'}) conf.initialize() # Register models for saving and loading. # `get_modules` gives a dictionary of `nn.Modules` in `conf`. # You can also specify a custom dictionary of models. experiment.add_pytorch_models(get_modules(conf)) # Start and watch the experiment with experiment.start(): # Run the training conf.run() def plot_image(img: torch.Tensor): """ ### Plots an image with matplotlib """ from matplotlib import pyplot as plt # Move tensor to CPU img = img.cpu() # Get min and max values of the image for normalization img_min, img_max = img.min(), img.max() # Scale image values to be [0...1] img = (img - img_min) / (img_max - img_min + 1e-5) # We have to change the order of dimensions to HWC. img = img.permute(1, 2, 0) # Show Image plt.imshow(img) # We don't need axes plt.axis('off') # Display plt.show() def evaluate(): """ ## Evaluate trained Cycle GAN """ # Set the run uuid from the training run trained_run_uuid = 'f73c1164184711eb9190b74249275441' # Create configs object conf = Configs() # Create experiment experiment.create(name='cycle_gan_inference') # Load hyper parameters set for training conf_dict = experiment.load_configs(trained_run_uuid) # Calculate configurations. We specify the generators `'generator_xy', 'generator_yx'` # so that it only loads those and their dependencies. # Configs like `device` and `img_channels` will be calculated since these are required by # `generator_xy` and `generator_yx`. # # If you want other parameters like `dataset_name` you should specify them here. # If you specify nothing all the configurations will be calculated including data loaders. # Calculation of configurations and their dependencies will happen when you call `experiment.start` experiment.configs(conf, conf_dict) conf.initialize() # Register models for saving and loading. # `get_modules` gives a dictionary of `nn.Modules` in `conf`. # You can also specify a custom dictionary of models. experiment.add_pytorch_models(get_modules(conf)) # Specify which run to load from. # Loading will actually happen when you call `experiment.start` experiment.load(trained_run_uuid) # Start the experiment with experiment.start(): # Image transformations transforms_ = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] # Load your own data, here we try test set. # I was trying with yosemite photos, they look awesome. # You can use `conf.dataset_name`, if you specified `dataset_name` as something you wanted to be calculated # in the call to `experiment.configs` dataset = ImageDataset(conf.dataset_name, transforms_, 'train') # Get an images from dataset x_image = dataset[10]['x'] # Display the image plot_image(x_image) # Evaluation mode conf.generator_xy.eval() conf.generator_yx.eval() # We dont need gradients with torch.no_grad(): # Add batch dimension and move to the device we use data = x_image.unsqueeze(0).to(conf.device) generated_y = conf.generator_xy(data) # Display the generated image. plot_image(generated_y[0].cpu()) if __name__ == '__main__': train() # evaluate()