def __init__( self, generator, adversary, encoder, x_dim, z_dim, optim=None, optim_kwargs=None, adv_type="Discriminator", feature_layer=None, fixed_noise_size=32, device=None, folder=None, ngpu=0, secure=True, _called_from_conditional=False): self.adv_type = adv_type self.generator = Generator(generator, input_size=z_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=x_dim, adv_type=adv_type, device=device, ngpu=ngpu, secure=secure) self.encoder = Encoder(encoder, input_size=x_dim, device=device, ngpu=ngpu, secure=secure) self.neural_nets = { "Generator": self.generator, "Adversary": self.adversary, "Encoder": self.encoder } super().__init__( x_dim=x_dim, z_dim=z_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, ngpu=ngpu, folder=folder, secure=secure ) self.hyperparameters["adv_type"] = adv_type if not _called_from_conditional and self.secure: assert self.generator.output_size == self.x_dim, ( "Generator output shape must be equal to x_dim. {} vs. {}.".format(self.generator.output_size, self.x_dim) )
def __init__( self, generator, adversary, encoder, x_dim, z_dim, y_dim, optim=None, optim_kwargs=None, adv_type="Discriminator", feature_layer=None, fixed_noise_size=32, device=None, folder="./veganModels/cAbstractConditionalGANGAE", ngpu=0, secure=True, _called_from_conditional=False): enc_in_dim = get_input_dim(dim1=x_dim, dim2=y_dim) gen_in_dim = get_input_dim(dim1=z_dim, dim2=y_dim) adv_in_dim = get_input_dim(dim1=x_dim, dim2=y_dim) if secure: AbstractConditionalGenerativeModel._check_conditional_network_input(encoder, in_dim=x_dim, y_dim=y_dim, name="Encoder") AbstractConditionalGenerativeModel._check_conditional_network_input(generator, in_dim=z_dim, y_dim=y_dim, name="Generator") AbstractConditionalGenerativeModel._check_conditional_network_input(adversary, in_dim=x_dim, y_dim=y_dim, name="Adversary") self.adv_type = adv_type self.encoder = Encoder(encoder, input_size=enc_in_dim, device=device, ngpu=ngpu, secure=secure) self.generator = Generator(generator, input_size=gen_in_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=adv_in_dim, adv_type=adv_type, device=device, ngpu=ngpu, secure=secure) self.neural_nets = { "Generator": self.generator, "Adversary": self.adversary, "Encoder": self.encoder } super().__init__( x_dim=x_dim, z_dim=z_dim, y_dim=y_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, ngpu=ngpu, folder=folder, secure=secure ) self.hyperparameters["adv_type"] = adv_type if not _called_from_conditional and self.secure: assert self.generator.output_size == self.x_dim, ( "Generator output shape must be equal to x_dim. {} vs. {}.".format(self.generator.output_size, self.x_dim) )
def __init__(self, generator, adversary, encoder, x_dim, z_dim, y_dim, c_dim_discrete, c_dim_continuous, optim=None, optim_kwargs=None, lambda_z=10, feature_layer=None, fixed_noise_size=32, device=None, ngpu=0, folder="./veganModels/cInfoGAN", secure=True): c_dim_discrete = [c_dim_discrete] if isinstance( c_dim_discrete, int) else c_dim_discrete assert c_dim_discrete == [0] or 0 not in c_dim_discrete, ( "`c_dim_discrete` has multiple elements. Zero not allowed. Given: {}." .format(c_dim_discrete)) assert isinstance( c_dim_continuous, int), ("`c_dim_continuous` must be of type int. Given: {}.".format( type(c_dim_continuous))) self.c_dim_discrete = tuple([i for i in list(c_dim_discrete)]) self.c_dim_continuous = tuple([c_dim_continuous]) self.c_dim = tuple( [sum(self.c_dim_discrete) + sum(self.c_dim_continuous)]) if len(y_dim) == 3: intermediate_in_dim = get_input_dim(dim1=z_dim, dim2=y_dim) gen_in_dim = get_input_dim(dim1=intermediate_in_dim, dim2=self.c_dim) else: gen_in_dim = get_input_dim(dim1=z_dim, dim2=sum(self.c_dim) + sum(y_dim)) adv_in_dim = get_input_dim(dim1=x_dim, dim2=y_dim) if secure: AbstractConditionalGenerativeModel._check_conditional_network_input( generator, in_dim=z_dim, y_dim=self.c_dim, name="Generator") self.generator = Generator(generator, input_size=gen_in_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=adv_in_dim, adv_type="Discriminator", device=device, ngpu=ngpu, secure=secure) self.encoder = Encoder(encoder, input_size=adv_in_dim, device=device, ngpu=ngpu, secure=secure) self.neural_nets = { "Generator": self.generator, "Adversary": self.adversary, "Encoder": self.encoder } super().__init__(x_dim=x_dim, z_dim=z_dim, y_dim=y_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, folder=folder, ngpu=ngpu, secure=secure) if self.c_dim_discrete != (0, ): self.multinomial = nn.Sequential( nn.Flatten(), nn.Linear(np.prod(self.encoder.output_size), np.sum(self.c_dim_discrete)), nn.Sigmoid()).to(self.device) if self.c_dim_continuous != (0, ): self.mu = nn.Sequential( nn.Flatten(), nn.Linear(np.prod(self.encoder.output_size), np.sum(self.c_dim_continuous)), LayerReshape(self.c_dim_continuous)).to(self.device) self.log_variance = nn.Sequential( nn.Flatten(), nn.Linear(np.prod(self.encoder.output_size), np.sum(self.c_dim_continuous)), nn.ReLU(), LayerReshape(self.c_dim_continuous)).to(self.device) self.lambda_z = lambda_z self.hyperparameters["lambda_z"] = lambda_z if self.secure: assert (self.generator.output_size == self.x_dim), ( "Generator output shape must be equal to x_dim. {} vs. {}.". format(self.generator.output_size, self.x_dim))
def __init__(self, generatorX_Y, adversaryX_Y, generatorY_X, adversaryY_X, x_dim, z_dim, y_dim, optim=None, optim_kwargs=None, lambda_x=10, adv_type="Discriminator", fixed_noise_size=32, device=None, ngpu=0, folder="./veganModels/cCycleGAN", secure=True): gen_in_dim = get_input_dim(dim1=z_dim, dim2=y_dim) adv_in_dim = get_input_dim(dim1=x_dim, dim2=y_dim) if secure: assert x_dim == y_dim, ( "`x_dim` and `y_dim` must be equal in the current implementation of CycleGAN. Given: {} and {}." .format(x_dim, y_dim)) AbstractConditionalGenerativeModel._check_conditional_network_input( generatorX_Y, in_dim=z_dim, y_dim=x_dim, name="GeneratorX_Y") AbstractConditionalGenerativeModel._check_conditional_network_input( adversaryX_Y, in_dim=y_dim, y_dim=x_dim, name="AdversaryX_Y") AbstractConditionalGenerativeModel._check_conditional_network_input( generatorY_X, in_dim=z_dim, y_dim=y_dim, name="GeneratorY_X") AbstractConditionalGenerativeModel._check_conditional_network_input( adversaryY_X, in_dim=x_dim, y_dim=y_dim, name="AdversaryY_X") self.adv_type = adv_type self.generatorX_Y = Generator(generatorX_Y, input_size=gen_in_dim, device=device, ngpu=ngpu, secure=secure) self.adversaryX_Y = Adversary(adversaryX_Y, input_size=adv_in_dim, adv_type=adv_type, device=device, ngpu=ngpu, secure=secure) self.generatorY_X = Generator(generatorY_X, input_size=gen_in_dim, device=device, ngpu=ngpu, secure=secure) self.adversaryY_X = Adversary(adversaryY_X, input_size=adv_in_dim, adv_type=adv_type, device=device, ngpu=ngpu, secure=secure) self.autoencoder = Autoencoder(encoder=self.generatorX_Y, decoder=self.generatorY_X) self.neural_nets = { "Autoencoder": self.autoencoder, "AdversaryX_Y": self.adversaryX_Y, "AdversaryY_X": self.adversaryY_X } self.generator = self.generatorX_Y self.adversary = self.adversaryX_Y super().__init__(x_dim=x_dim, z_dim=z_dim, y_dim=y_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=None, fixed_noise_size=fixed_noise_size, device=device, folder=folder, ngpu=ngpu, secure=secure) self.lambda_x = lambda_x self.hyperparameters["lambda_x"] = lambda_x if self.secure: assert (self.generatorX_Y.output_size == self.x_dim), ( "GeneratorX_Y output shape must be equal to x_dim. {} vs. {}.". format(self.generatorX_Y.output_size, self.x_dim)) assert (self.generatorY_X.output_size == self.x_dim), ( "GeneratorY_X output shape must be equal to x_dim. {} vs. {}.". format(self.generatorY_X.output_size, self.x_dim))
class AbstractGANGAE(AbstractGenerativeModel): """ Abstract class for GAN with structure of one generator, one discriminator / critic and one encoder. Examples are the `LRGAN`, `VAEGAN` and `BicycleGAN`. Parameters ---------- generator: nn.Module Generator architecture. Produces output in the real space. adversary: nn.Module Adversary architecture. Produces predictions for real and fake samples to differentiate them. encoder : nn.Module Encoder architecture. Produces predictions in the latent space. x_dim : list, tuple Number of the output dimensions of the generator and input dimension of the discriminator / critic. In the case of images this will be [nr_channels, nr_height_pixels, nr_width_pixels]. z_dim : int, list, tuple Number of the latent dimensions for the generator input. Might have dimensions of an image. optim : dict or torch.optim Optimizer used for each network. Could be either an optimizer from torch.optim or a dictionary with network name keys and torch.optim as value, i.e. {"Generator": torch.optim.Adam}. optim_kwargs : dict Optimizer keyword arguments used for each network. Must be a dictionary with network name keys and dictionary with keyword arguments as value, i.e. {"Generator": {"lr": 0.0001}}. feature_layer : torch.nn.* Output layer used to compute the feature loss. Should be from either the discriminator or critic. If `feature_layer` is not None, the original generator loss is replaced by a feature loss, introduced [here](https://arxiv.org/abs/1606.03498v1). fixed_noise_size : int Number of images shown when logging. The fixed noise is used to produce the images in the folder/images subdirectory, the tensorboard images tab and the samples in get_training_results(). lambda_grad: float Weight for the reconstruction loss of the gradients. Pushes the norm of the gradients to 1. device : string Device used while training the model. Either "cpu" or "cuda". ngpu : int Number of gpus used during training if device == "cuda". folder : string Creates a folder in the current working directory with this name. All relevant files like summary, images, models and tensorboard output are written there. Existing folders are never overwritten or deleted. If a folder with the same name already exists a time stamp is appended to make it unique. """ ######################################################################### # Actions before training ######################################################################### def __init__( self, generator, adversary, encoder, x_dim, z_dim, optim=None, optim_kwargs=None, adv_type="Discriminator", feature_layer=None, fixed_noise_size=32, device=None, folder=None, ngpu=0, secure=True, _called_from_conditional=False): self.adv_type = adv_type self.generator = Generator(generator, input_size=z_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=x_dim, adv_type=adv_type, device=device, ngpu=ngpu, secure=secure) self.encoder = Encoder(encoder, input_size=x_dim, device=device, ngpu=ngpu, secure=secure) self.neural_nets = { "Generator": self.generator, "Adversary": self.adversary, "Encoder": self.encoder } super().__init__( x_dim=x_dim, z_dim=z_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, ngpu=ngpu, folder=folder, secure=secure ) self.hyperparameters["adv_type"] = adv_type if not _called_from_conditional and self.secure: assert self.generator.output_size == self.x_dim, ( "Generator output shape must be equal to x_dim. {} vs. {}.".format(self.generator.output_size, self.x_dim) ) def _define_loss(self): if self.adv_type == "Discriminator": loss_functions = {"Generator": BCELoss(), "Adversary": BCELoss()} elif self.adv_type == "Critic": loss_functions = {"Generator": WassersteinLoss(), "Adversary": WassersteinLoss()} else: raise NotImplementedError("'adv_type' must be one of Discriminator or Critic.") return loss_functions ######################################################################### # Actions during training ######################################################################### def calculate_losses(self, X_batch, Z_batch, who=None): """ Calculates the losses for GANs using a 1v1 architecture. This method is called within the `AbstractGenerativeModel` main `fit()` loop. Parameters ---------- X_batch : torch.Tensor Current x batch. Z_batch : torch.Tensor Current z batch. who : None, optional Name of the network that should be trained. """ if who == "Generator": losses = self._calculate_generator_loss(X_batch=X_batch, Z_batch=Z_batch) elif who == "Adversary": losses = self._calculate_adversary_loss(X_batch=X_batch, Z_batch=Z_batch) elif who == "Encoder": losses = self._calculate_encoder_loss(X_batch=X_batch, Z_batch=Z_batch) else: losses = self._calculate_generator_loss(X_batch=X_batch, Z_batch=Z_batch) losses.update(self._calculate_adversary_loss(X_batch=X_batch, Z_batch=Z_batch)) losses.update(self._calculate_encoder_loss(X_batch=X_batch, Z_batch=Z_batch)) return losses def _step(self, who=None): if who is not None: self.optimizers[who].step() if who == "Adversary": if self.adv_type == "Critic": for p in self.adversary.parameters(): p.data.clamp_(-0.01, 0.01) else: [optimizer.step() for _, optimizer in self.optimizers.items()] ######################################################################### # Utility functions ######################################################################### def encode(self, x): return self.encoder(x)
def __init__(self, generator, adversary, encoder, x_dim, z_dim, optim=None, optim_kwargs=None, lambda_z=10, adv_type="Discriminator", feature_layer=None, fixed_noise_size=32, device=None, ngpu=0, folder="./veganModels/AAE", secure=True): self.adv_type = adv_type self.encoder = Encoder(encoder, input_size=x_dim, device=device, ngpu=ngpu, secure=secure) self.generator = Generator(generator, input_size=z_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=z_dim, device=device, ngpu=ngpu, adv_type=adv_type, secure=secure) self.neural_nets = { "Generator": self.generator, "Encoder": self.encoder, "Adversary": self.adversary } super().__init__(x_dim=x_dim, z_dim=z_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, folder=folder, ngpu=ngpu, secure=secure) self.lambda_z = lambda_z self.hyperparameters["lambda_z"] = lambda_z self.hyperparameters["adv_type"] = adv_type if self.secure: assert self.encoder.output_size == self.z_dim, ( "Encoder output shape must be equal to z_dim. {} vs. {}.". format(self.encoder.output_size, self.z_dim)) assert self.generator.output_size == self.x_dim, ( "Generator output shape must be equal to x_dim. {} vs. {}.". format(self.generator.output_size, self.x_dim))
class AAE(AbstractGenerativeModel): """ Parameters ---------- generator: nn.Module Generator architecture. Produces output in the real space. adversary: nn.Module Adversary architecture. Produces predictions for real and fake samples to differentiate them. encoder: nn.Module Encoder architecture. Produces predictions in the latent space. x_dim : list, tuple Number of the output dimensions of the generator and input dimension of the discriminator / critic. In the case of images this will be [nr_channels, nr_height_pixels, nr_width_pixels]. z_dim : int, list, tuple Number of the latent dimensions for the generator input. Might have dimensions of an image. optim : dict or torch.optim Optimizer used for each network. Could be either an optimizer from torch.optim or a dictionary with network name keys and torch.optim as value, i.e. {"Generator": torch.optim.Adam}. optim_kwargs : dict Optimizer keyword arguments used for each network. Must be a dictionary with network name keys and dictionary with keyword arguments as value, i.e. {"Generator": {"lr": 0.0001}}. lambda_z: float Weight for the discriminator loss computing the realness of the latent z dimension. adv_type: "Discriminator", "Critic" or "Autoencoder" Indicating which adversarial architecture will be used. feature_layer : torch.nn.* Output layer used to compute the feature loss. Should be from either the discriminator or critic. If `feature_layer` is not None, the original generator loss is replaced by a feature loss, introduced [here](https://arxiv.org/abs/1606.03498v1). fixed_noise_size : int Number of images shown when logging. The fixed noise is used to produce the images in the folder/images subdirectory, the tensorboard images tab and the samples in get_training_results(). device : string Device used while training the model. Either "cpu" or "cuda". ngpu : int Number of gpus used during training if device == "cuda". folder : string Creates a folder in the current working directory with this name. All relevant files like summary, images, models and tensorboard output are written there. Existing folders are never overwritten or deleted. If a folder with the same name already exists a time stamp is appended to make it unique. """ ######################################################################### # Actions before training ######################################################################### def __init__(self, generator, adversary, encoder, x_dim, z_dim, optim=None, optim_kwargs=None, lambda_z=10, adv_type="Discriminator", feature_layer=None, fixed_noise_size=32, device=None, ngpu=0, folder="./veganModels/AAE", secure=True): self.adv_type = adv_type self.encoder = Encoder(encoder, input_size=x_dim, device=device, ngpu=ngpu, secure=secure) self.generator = Generator(generator, input_size=z_dim, device=device, ngpu=ngpu, secure=secure) self.adversary = Adversary(adversary, input_size=z_dim, device=device, ngpu=ngpu, adv_type=adv_type, secure=secure) self.neural_nets = { "Generator": self.generator, "Encoder": self.encoder, "Adversary": self.adversary } super().__init__(x_dim=x_dim, z_dim=z_dim, optim=optim, optim_kwargs=optim_kwargs, feature_layer=feature_layer, fixed_noise_size=fixed_noise_size, device=device, folder=folder, ngpu=ngpu, secure=secure) self.lambda_z = lambda_z self.hyperparameters["lambda_z"] = lambda_z self.hyperparameters["adv_type"] = adv_type if self.secure: assert self.encoder.output_size == self.z_dim, ( "Encoder output shape must be equal to z_dim. {} vs. {}.". format(self.encoder.output_size, self.z_dim)) assert self.generator.output_size == self.x_dim, ( "Generator output shape must be equal to x_dim. {} vs. {}.". format(self.generator.output_size, self.x_dim)) def _define_loss(self): if self.adv_type == "Discriminator": loss_functions = {"Generator": MSELoss(), "Adversary": BCELoss()} elif self.adv_type == "Critic": loss_functions = { "Generator": MSELoss(), "Adversary": WassersteinLoss() } else: raise NotImplementedError( "'adv_type' must be one of Discriminator or Critic.") return loss_functions ######################################################################### # Actions during training ######################################################################### def encode(self, x): return self.encoder(x) def calculate_losses(self, X_batch, Z_batch, who=None): if who == "Generator": losses = self._calculate_generator_loss(X_batch=X_batch, Z_batch=Z_batch) elif who == "Encoder": losses = self._calculate_encoder_loss(X_batch=X_batch, Z_batch=Z_batch) elif who == "Adversary": losses = self._calculate_adversary_loss(X_batch=X_batch, Z_batch=Z_batch) else: losses = self._calculate_generator_loss(X_batch=X_batch, Z_batch=Z_batch) losses.update( self._calculate_encoder_loss(X_batch=X_batch, Z_batch=Z_batch)) losses.update( self._calculate_adversary_loss(X_batch=X_batch, Z_batch=Z_batch)) return losses def _calculate_generator_loss(self, X_batch, Z_batch, fake_images=None): if fake_images is None: encoded_output = self.encode(x=X_batch).detach() fake_images = self.generate(encoded_output) gen_loss = self.loss_functions["Generator"](fake_images, X_batch) return { "Generator": gen_loss, } def _calculate_encoder_loss(self, X_batch, Z_batch, fake_images=None, encoded_output=None): if fake_images is None: encoded_output = self.encode(x=X_batch) fake_images = self.generate(z=encoded_output) if self.feature_layer is None: fake_predictions = self.predict(x=encoded_output) enc_loss_fake = self.loss_functions["Generator"]( fake_predictions, torch.ones_like(fake_predictions, requires_grad=False)) else: enc_loss_fake = self._calculate_feature_loss(X_real=Z_batch, X_fake=encoded_output) enc_loss_reconstruction = self.loss_functions["Generator"](fake_images, X_batch) enc_loss = self.lambda_z * enc_loss_fake + enc_loss_reconstruction return { "Encoder": enc_loss, "Encoder_x": self.lambda_z * enc_loss_fake, "Encoder_fake": enc_loss_reconstruction, } def _calculate_adversary_loss(self, X_batch, Z_batch, encoded_output=None): if encoded_output is None: encoded_output = self.encode(x=X_batch).detach() fake_predictions = self.predict(x=encoded_output) real_predictions = self.predict(x=Z_batch) adv_loss_fake = self.loss_functions["Adversary"]( fake_predictions, torch.zeros_like(fake_predictions, requires_grad=False)) adv_loss_real = self.loss_functions["Adversary"]( real_predictions, torch.ones_like(real_predictions, requires_grad=False)) adv_loss = 1 / 2 * (adv_loss_real + adv_loss_fake) return { "Adversary": adv_loss, "Adversary_fake": adv_loss_fake, "Adversary_real": adv_loss_real, "RealFakeRatio": adv_loss_real / adv_loss_fake } def _step(self, who=None): if who is not None: self.optimizers[who].step() if who == "Adversary": if self.adv_type == "Critic": for p in self.adversary.parameters(): p.data.clamp_(-0.01, 0.01) else: [optimizer.step() for _, optimizer in self.optimizers.items()]