def __init__(self, channel_num: int, z_dim: int, alpha: float, beta: float, gamma: float, **kwargs): super().__init__() # Parameters self.channel_num = channel_num self.z_dim = z_dim self._alpha_value = alpha self._beta_value = beta self._gamma_value = gamma # Distributions self.prior = pxd.Normal( loc=torch.zeros(z_dim), scale=torch.ones(z_dim), var=["z"]) self.decoder = Decoder(channel_num, z_dim) self.encoder = Encoder(channel_num, z_dim) self.distributions = [self.prior, self.decoder, self.encoder] # Loss class self.ce = pxl.CrossEntropy(self.encoder, self.decoder) self.kl = pxl.KullbackLeibler(self.encoder, self.prior) self.alpha = pxl.Parameter("alpha") self.beta = pxl.Parameter("beta") self.gamma = pxl.Parameter("gamma")
def __init__(self, channel_num: int, z_dim: int, beta: float, c: float, lmd_od: float, lmd_d: float, dip_type: str, **kwargs): super().__init__() # Parameters self.channel_num = channel_num self.z_dim = z_dim self._beta_value = beta self._c_value = c self.lmd_od = lmd_od self.lmd_d = lmd_d # Distributions self.prior = pxd.Normal(loc=torch.zeros(z_dim), scale=torch.ones(z_dim), var=["z"]) self.decoder = Decoder(channel_num, z_dim) self.encoder = Encoder(channel_num, z_dim) self.distributions = [self.prior, self.decoder, self.encoder] # Loss class self.ce = pxl.CrossEntropy(self.encoder, self.decoder) _kl = pxl.KullbackLeibler(self.encoder, self.prior) _beta = pxl.Parameter("beta") _c = pxl.Parameter("c") self.kl = _beta * (_kl - _c).abs() self.dip = DipLoss(self.encoder, lmd_od, lmd_d, dip_type)
def load_dmm_model(x_dim, t_max, device, args): # Latent dimensions h_dim = args.h_dim hidden_dim = args.hidden_dim z_dim = args.z_dim # Distributions prior = Prior(z_dim, hidden_dim).to(device) decoder = Generator(z_dim, hidden_dim, x_dim).to(device) encoder = Inference(z_dim, h_dim).to(device) rnn = RNN(x_dim, h_dim).to(device) # Sampler generate_from_prior = prior * decoder # Loss ce = pxl.CrossEntropy(encoder, decoder) kl = pxl.KullbackLeibler(encoder, prior) step_loss = ce + kl _loss = pxl.IterativeLoss(step_loss, max_iter=t_max, series_var=["x", "h"], update_value={"z": "z_prev"}) loss = _loss.expectation(rnn).mean() # Model dmm = pxm.Model(loss, distributions=[rnn, encoder, decoder, prior], optimizer=optim.Adam, optimizer_params={"lr": args.learning_rate, "betas": (args.beta1, args.beta2), "weight_decay": args.weight_decay}, clip_grad_norm=args.clip_grad_norm) return dmm, generate_from_prior, decoder
def __init__(self, x_dim, z_dim, h_dim): # Generative model self.prior = pxd.Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["z"], features_shape=[z_dim]) self.decoder = Generator(z_dim, h_dim, x_dim) # Variational model self.encoder = Inference(x_dim, h_dim, z_dim) # Loss ce = pxl.CrossEntropy(self.encoder, self.decoder) kl = pxl.KullbackLeibler(self.encoder, self.prior) loss = (ce + kl).mean() # Init super().__init__(loss, distributions=[self.encoder, self.decoder])
def __init__(self, channel_num: int, z_dim: int, c_dim: int, temperature: float, gamma_z: float, gamma_c: float, cap_z: float, cap_c: float, **kwargs): super().__init__() self.channel_num = channel_num self.z_dim = z_dim self.c_dim = c_dim self._gamma_z_value = gamma_z self._gamma_c_value = gamma_c self._cap_z_value = cap_z self._cap_c_value = cap_c # Distributions self.prior_z = pxd.Normal( loc=torch.zeros(z_dim), scale=torch.ones(z_dim), var=["z"]) self.prior_c = pxd.Categorical( probs=torch.ones(c_dim, dtype=torch.float32) / c_dim, var=["c"]) self.encoder_func = EncoderFunction(channel_num) self.encoder_z = ContinuousEncoder(z_dim) self.encoder_c = DiscreteEncoder(c_dim, temperature) self.decoder = JointDecoder(channel_num, z_dim, c_dim) self.distributions = [self.prior_z, self.prior_c, self.encoder_func, self.encoder_z, self.encoder_c, self.decoder] # Loss self.ce = pxl.CrossEntropy(self.encoder_z * self.encoder_c, self.decoder) self.kl_z = pxl.KullbackLeibler(self.encoder_z, self.prior_z) self.kl_c = CategoricalKullbackLeibler( self.encoder_c, self.prior_c) # Coefficient for kl self.gamma_z = pxl.Parameter("gamma_z") self.gamma_c = pxl.Parameter("gamma_c") # Capacity self.cap_z = pxl.Parameter("cap_z") self.cap_c = pxl.Parameter("cap_c")
def main(): # ------------------------------------------------------------------------- # 1. Settings # ------------------------------------------------------------------------- # Args args = init_args() # Settings use_cuda = args.cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") torch.manual_seed(args.seed) # Tensorboard writer writer = tensorboard.SummaryWriter(args.logdir) # ------------------------------------------------------------------------- # 2. Data # ------------------------------------------------------------------------- # Loader batch_size = args.batch_size train_loader, test_loader = init_dataloader(root=args.data_root, cuda=use_cuda, batch_size=batch_size) # Data dimension x_dim = train_loader.dataset.data.shape[1] t_max = train_loader.dataset.data.shape[2] # ------------------------------------------------------------------------- # 3. Model # ------------------------------------------------------------------------- # Latent dimension h_dim = args.h_dim hidden_dim = args.hidden_dim z_dim = args.z_dim # Distributions prior = Prior(z_dim, hidden_dim).to(device) decoder = Generator(z_dim, hidden_dim, x_dim).to(device) encoder = Inference(z_dim, h_dim).to(device) rnn = RNN(x_dim, h_dim).to(device) # Sampler generate_from_prior = prior * decoder # Loss ce = pxl.CrossEntropy(encoder, decoder) kl = pxl.KullbackLeibler(encoder, prior) _loss = pxl.IterativeLoss(ce + kl, max_iter=t_max, series_var=["x", "h"], update_value={"z": "z_prev"}) loss = _loss.expectation(rnn).mean() # Model model = pxm.Model(loss, distributions=[rnn, encoder, decoder, prior], optimizer=optim.Adam, optimizer_params={"lr": 1e-3}, clip_grad_value=10) # ------------------------------------------------------------------------- # 4. Training # ------------------------------------------------------------------------- for epoch in range(1, args.epochs + 1): # Training train_loss = data_loop(train_loader, model, z_dim, device, train_mode=True) test_loss = data_loop(test_loader, model, z_dim, device, train_mode=False) # Sample data sample = plot_image_from_latent(generate_from_prior, decoder, batch_size, z_dim, t_max, device) # Log writer.add_scalar("train_loss", train_loss.item(), epoch) writer.add_scalar("test_loss", test_loss.item(), epoch) writer.add_images("image_from_latent", sample, epoch) writer.close()
def main(): # ------------------------------------------------------------------------- # 1. Settings # ------------------------------------------------------------------------- # Args args = init_args() # Settings use_cuda = args.cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") torch.manual_seed(args.seed) # Tensorboard writer writer = tensorboard.SummaryWriter(args.logdir) # ------------------------------------------------------------------------- # 2. Data # ------------------------------------------------------------------------- # Loader batch_size = args.batch_size train_loader, test_loader = init_dataloader(root=args.data_root, cuda=use_cuda, batch_size=batch_size) # Data dimension x_dim = train_loader.dataset.data.shape[1] t_max = train_loader.dataset.data.shape[2] # ------------------------------------------------------------------------- # 3. Model # ------------------------------------------------------------------------- # Latent dimension h_dim = args.h_dim z_dim = args.z_dim # Functions f_phi_x = Phi_x(x_dim, h_dim).to(device) f_phi_z = Phi_z(z_dim, h_dim).to(device) # Distributions prior = Prior(h_dim, z_dim).to(device) decoder = Generator(h_dim, z_dim, x_dim, f_phi_z).to(device) encoder = Inference(h_dim, z_dim, f_phi_x).to(device) recurrence = Recurrence(h_dim, f_phi_x, f_phi_z).to(device) # Sampler encoder_with_recurrence = encoder * recurrence generate_from_prior = prior * decoder * recurrence # Loss reconst = pxl.StochasticReconstructionLoss(encoder_with_recurrence, decoder) kl = pxl.KullbackLeibler(encoder, prior) step_loss = (reconst + kl).mean() loss = pxl.IterativeLoss(step_loss, max_iter=t_max, series_var=["x"], update_value={"h": "h_prev"}) # Model vrnn = pxm.Model(loss, distributions=[encoder, decoder, prior, recurrence], optimizer=optim.Adam, optimizer_params={"lr": 1e-3}) # ------------------------------------------------------------------------- # 4. Training # ------------------------------------------------------------------------- for epoch in range(1, args.epochs + 1): # Training train_loss = data_loop(train_loader, vrnn, h_dim, device, train_mode=True) test_loss = data_loop(test_loader, vrnn, h_dim, device, train_mode=False) # Sample data sample = plot_image_from_latent(generate_from_prior, decoder, batch_size, h_dim, t_max, device) # Log writer.add_scalar("train_loss", train_loss.item(), epoch) writer.add_scalar("test_loss", test_loss.item(), epoch) writer.add_images("image_from_latent", sample, epoch) writer.close()
def main(): # ------------------------------------------------------------------------- # 1. Settings # ------------------------------------------------------------------------- # Args args = init_args() # Settings use_cuda = args.cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") torch.manual_seed(args.seed) # Tensorboard writer writer = tensorboard.SummaryWriter(args.logdir) # ------------------------------------------------------------------------- # 2. Data # ------------------------------------------------------------------------- # Loader batch_size = args.batch_size train_loader, test_loader = init_dataloader( root=args.data_root, cuda=use_cuda, batch_size=batch_size) # Sample data _x, _ = iter(test_loader).next() _x = _x.to(device) # Data dimension x_dim = _x.size(1) image_dim = int(x_dim ** 0.5) # Latent dimension z_dim = args.z_dim # Dummy latent variable z_sample = 0.5 * torch.randn(args.plot_num, z_dim).to(device) # ------------------------------------------------------------------------- # 3. Pixyz classses # ------------------------------------------------------------------------- # Distributions p = Generator(z_dim, x_dim).to(device) q = Inference(x_dim, z_dim).to(device) d = Discriminator(z_dim).to(device) q_shuffle = InferenceShuffleDim(q).to(device) prior = pxd.Normal( loc=torch.tensor(0.), scale=torch.tensor(1.), var=["z"], features_shape=[z_dim], name="p_prior").to(device) # Loss reconst = -q.log_prob().expectation(q) kl = pxl.KullbackLeibler(q, prior) tc = pxl.AdversarialKullbackLeibler(q, q_shuffle, d, optimizer=optim.Adam, optimizer_params={"lr": 1e-3}) loss_cls = reconst.mean() + kl.mean() + 10 * tc # Model model = pxm.Model(loss_cls, distributions=[p, q], optimizer=optim.Adam, optimizer_params={"lr": 1e-3}) # ------------------------------------------------------------------------- # 4. Training # ------------------------------------------------------------------------- for epoch in range(1, args.epochs + 1): # Training train_loss, train_d_loss = data_loop(train_loader, model, tc, device, train_mode=True) test_loss, test_d_loss = data_loop(test_loader, model, tc, device, train_mode=False) # Sample data recon = plot_reconstruction( _x[:args.plot_recon], q, p, image_dim, image_dim) sample = plot_image_from_latent(z_sample, p, image_dim, image_dim) # Log writer.add_scalar("train_loss", train_loss.item(), epoch) writer.add_scalar("test_loss", test_loss.item(), epoch) writer.add_scalar("train_d_loss", train_d_loss.item(), epoch) writer.add_scalar("test_d_loss", test_d_loss.item(), epoch) writer.add_images("image_reconstruction", recon, epoch) writer.add_images("image_from_latent", sample, epoch) writer.close()
def main(): # ------------------------------------------------------------------------- # 1. Settings # ------------------------------------------------------------------------- # Args args = init_args() # Settings use_cuda = args.cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") torch.manual_seed(args.seed) # Tensorboard writer writer = tensorboard.SummaryWriter(args.logdir) # ------------------------------------------------------------------------- # 2. Data # ------------------------------------------------------------------------- # Loader train_loader, test_loader = init_dataloader(root=args.data_root, cuda=use_cuda, batch_size=args.batch_size) # Sample data for comparison _x, _ = iter(test_loader).next() _x = _x.to(device) # Data dimension x_dim = _x.shape[1] image_dim = int(x_dim**0.5) # Latent dimension z_dim = args.z_dim # Latent data for visualization z_sample = 0.5 * torch.randn(args.plot_dim, z_dim).to(device) # ------------------------------------------------------------------------- # 3. Model # ------------------------------------------------------------------------- # Distributions p = Generator(z_dim, x_dim).to(device) q = Inference(x_dim, z_dim).to(device) prior = pxd.Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["z"], name="p_prior").to(device) # Model model = pxm.VAE(q, p, regularizer=pxl.KullbackLeibler(q, prior), optimizer=optim.Adam, optimizer_params={"lr": 1e-3}) # ------------------------------------------------------------------------- # 4. Training # ------------------------------------------------------------------------- for epoch in range(1, args.epochs + 1): train_loss = data_loop(train_loader, model, device, train_mode=True) test_loss = data_loop(test_loader, model, device, train_mode=False) recon = plot_reconstruction(_x[:args.plot_recon], q, p, image_dim, image_dim) sample = plot_image_from_latent(z_sample, p, image_dim, image_dim) writer.add_scalar("train_loss", train_loss.item(), epoch) writer.add_scalar("test_loss", test_loss.item(), epoch) writer.add_images("image_reconstruction", recon, epoch) writer.add_images("image_from_latent", sample, epoch) writer.close()