def __init__( self, generator_and_opt: [Generator, torch.optim] = None, discriminator_and_opt: [Discriminator, torch.optim] = None, input_size: int = None, hidden_channel: int = 128, latent_dim: int = 100, learning_rate: float = 1e-4, ): self.generator = None self.discriminator = None super().__init__() # Generator if generator_and_opt is None: assert input_size is None, "generator_and_opt or input_size should be given." if self.generator is not None: self.generator = Generator(input_size=input_size, latent_dim=latent_dim, hidden_channel=hidden_channel) self.generator_opt = torch.optim.Adam( self.generator.parameters(), learning_rate) else: self.generator, self.generator_opt = generator_and_opt # Discriminator if discriminator_and_opt is None: assert input_size is None, "discriminator_and_opt or input_size should be given." if self.discriminator is not None: self.discriminator = Discriminator( input_size=input_size, hidden_channel=hidden_channel) self.discriminator_opt = torch.optim.Adam( self.discriminator.parameters(), learning_rate) else: self.discriminator, self.discriminator_opt = discriminator_and_opt
def build_model(self): self.G = Generator() self.D = Discriminator() self.G_optim = optim.Adam(self.G.parameters(), self.G_lr, (self.beta1, 0.999)) self.D_optim = optim.Adam(self.D.parameters(), self.D_lr, (self.beta1, 0.999)) if self.loss_type == 'BCEwL': self.criterion = nn.BCEWithLogitsLoss() elif self.loss_type == 'WGAN': pass elif self.loss_type == 'WGAN+': pass self.fixed_sample = None self.fixed_noise = None # self.true = torch.ones([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device) # self.false = torch.zeros([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device) # Change to GPU mode print('Change CPU mode to GPU mode...') self.G.to(self.device) self.D.to(self.device) print('Creating models are success...')
def __init__(self, args): self.args = args self.pretrained = False self.epoch = 0 self.G = Generator() self.D = Discriminator() self.g_optimizer = optim.Adam(self.G.parameters(), lr=1E-4) self.d_optimizer = optim.Adam(self.D.parameters(), lr=1E-4) self.g_scheduler = optim.lr_scheduler.StepLR(self.g_optimizer, step_size=40) self.d_scheduler = optim.lr_scheduler.StepLR(self.d_optimizer, step_size=40) self.train_losses = [] self.val_losses = [] if args.load_model: self._load_state(args.load_model) # extract all layers prior to the last softmax of VGG-19 vgg19_layers = list(models.vgg19(pretrained=True).features)[:36] self.vgg19 = nn.Sequential(*vgg19_layers).eval() for param in self.vgg19.parameters(): param.requires_grad = False self.mse_loss = torch.nn.MSELoss() self.bce_loss = torch.nn.BCELoss()
def val_test(args): writer = SummaryWriter('./logs/{0}'.format(args.output_folder)) save_filename = './models/{0}'.format(args.output_folder) train_loader, valid_loader, test_loader = train_util.get_dataloaders(args) recons_input_img = train_util.log_input_img_grid(test_loader, writer) input_dim = 3 model = VectorQuantizedVAE(input_dim, args.hidden_size, args.k, args.enc_type, args.dec_type) # if torch.cuda.device_count() > 1 and args.device == "cuda": # model = torch.nn.DataParallel(model) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) discriminators = {} if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] model.to(args.device) for disc in discriminators: discriminators[disc][0].to(args.device) if args.weights == "load": start_epoch = train_util.load_state(save_filename, model, optimizer, discriminators) else: start_epoch = 0 stop_patience = args.stop_patience best_loss = torch.tensor(np.inf) for epoch in tqdm(range(start_epoch, 4), file=sys.stdout): val_loss_dict, z = train_util.test(get_losses, model, valid_loader, args, discriminators, True) # if args.weights == "init" and epoch==1: # epoch+=1 # break train_util.log_recons_img_grid(recons_input_img, model, epoch + 1, args.device, writer) train_util.log_interp_img_grid(recons_input_img, model, epoch + 1, args.device, writer) train_util.log_losses("val", val_loss_dict, epoch + 1, writer) train_util.log_latent_metrics("val", z, epoch + 1, writer) train_util.save_state(model, optimizer, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, save_filename) print(val_loss_dict)
def build_model(self): self.G = Generator(image_size=self.imsize, z_dim=self.z_dim, conv_dim=self.g_dim) self.D = Discriminator(conv_dim=self.d_dim) self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2]) self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2]) self.print_network(self.G, 'G') self.print_network(self.D, 'D') self.G.to(self.device) self.D.to(self.device)
def build_models(self): self.E = Encoder(self.c64, self.rb6) self.T_Hair = Transformer(self.hair_dim, self.c256, self.rb6) self.T_Gender = Transformer(self.attr_dim, self.c256, self.rb6) self.T_Smailing = Transformer(self.attr_dim, self.c256, self.rb6) self.R = Reconstructor(self.c256) self.D_Hair = Discriminator(self.hair_dim, self.c64) self.D_Gender = Discriminator(self.attr_dim, self.c64) self.D_Smailing = Discriminator(self.attr_dim, self.c64) self.e_optim = torch.optim.Adam(self.E.parameters(), self.e_lr, [self.beta1, self.beta2]) self.th_optim = torch.optim.Adam(self.T_Hair.parameters(), self.t_lr, [self.beta1, self.beta2]) self.tg_optim = torch.optim.Adam(self.T_Gender.parameters(), self.t_lr, [self.beta1, self.beta2]) self.ts_optim = torch.optim.Adam(self.T_Smailing.parameters(), self.t_lr, [self.beta1, self.beta2]) self.r_optim = torch.optim.Adam(self.R.parameters(), self.r_lr, [self.beta1, self.beta2]) self.dh_optim = torch.optim.Adam(self.D_Hair.parameters(), self.d_lr, [self.beta1, self.beta2]) self.dg_optim = torch.optim.Adam(self.D_Gender.parameters(), self.d_lr, [self.beta1, self.beta2]) self.ds_optim = torch.optim.Adam(self.D_Smailing.parameters(), self.d_lr, [self.beta1, self.beta2]) self.print_network(self.E, 'Encoder') self.print_network(self.T_Hair, 'Transformer for Hair Color') self.print_network(self.T_Gender, 'Transformer for Gender') self.print_network(self.T_Smailing, 'Transformer for Smailing') self.print_network(self.R, 'Reconstructor') self.print_network(self.D_Hair, 'D for Hair Color') self.print_network(self.D_Gender, 'D for Gender') self.print_network(self.D_Smailing, 'D for Smailing') self.E.to(self.device) self.T_Hair.to(self.device) self.T_Gender.to(self.device) self.T_Smailing.to(self.device) self.R.to(self.device) self.D_Gender.to(self.device) self.D_Smailing.to(self.device) self.D_Hair.to(self.device)
def train(dataset_dir, output_dir): """Train discriminator and generator""" if torch.cuda.is_available(): torch.set_default_tensor_type('torch.cuda.FloatTensor') device = torch.device('cuda') else: device = torch.device('cpu') discriminator = Discriminator() discriminator.to(device) discriminator.train() generator = Generator() generator.to(device) generator.train() optimizer_discriminator = torch.optim.Adam(discriminator.parameters(), lr=0.0001) optimizer_generator = torch.optim.Adam(generator.parameters(), lr=0.0001) loss_func = nn.BCELoss() loader = LFWLoader(dataset_dir) if not os.path.exists(output_dir): os.makedirs(output_dir) for epoch in range(10): for image in tqdm(loader): # train discriminator on true outputs = discriminator(torch.tensor(image, dtype=torch.float32) / 255) loss = loss_func(outputs, torch.tensor([1.0])) optimizer_discriminator.zero_grad() loss.backward() optimizer_discriminator.step() # train discriminator on false outputs = discriminator(generator(generate_random_seed(100)).detach()) loss = loss_func(outputs, torch.tensor([0.0])) optimizer_discriminator.zero_grad() loss.backward() optimizer_discriminator.step() # train generator outputs = discriminator(generator(generate_random_seed(100))) loss = loss_func(outputs, torch.tensor([1.0])) optimizer_generator.zero_grad() loss.backward() optimizer_generator.step() torch.save(generator, os.path.join(output_dir, 'G_%d.pt' % epoch)) torch.save(discriminator, os.path.join(output_dir, 'D_%d.pt' % epoch))
def __init__(self, N_FEATURES=1, N_CLASSES=2, HIDDEN_DIM=50, CELL_TYPE="LSTM", N_LAYERS=1, DF=1., LAMBDA=1.0): super(EARLIEST, self).__init__() # --- Hyperparameters --- self.CELL_TYPE = CELL_TYPE self.HIDDEN_DIM = HIDDEN_DIM self.DF = DF self.LAMBDA = torch.tensor([LAMBDA], requires_grad=False) self.N_LAYERS = N_LAYERS self._epsilon = 1.0 self._rewards = 0 # --- Sub-networks --- self.BaseRNN = BaseRNN(N_FEATURES, HIDDEN_DIM, CELL_TYPE) self.Controller = Controller(HIDDEN_DIM+1, 1) # Add +1 for timestep input self.BaselineNetwork = BaselineNetwork(HIDDEN_DIM, 1) self.Discriminator = Discriminator(HIDDEN_DIM, N_CLASSES)
def __init__(self,config): super(BiGAN,self).__init__() self._work_type = config.work_type self._epochs = config.epochs self._batch_size = config.batch_size self._encoder_lr = config.encoder_lr self._generator_lr = config.generator_lr self._discriminator_lr = config.discriminator_lr self._latent_dim = config.latent_dim self._weight_decay = config.weight_decay self._img_shape = (config.input_size,config.input_size) self._img_save_path = config.image_save_path self._model_save_path = config.model_save_path self._device = config.device if self._work_type == 'train': # Loss function self._adversarial_criterion = torch.nn.MSELoss() # Initialize generator, encoder and discriminator self._G = Generator(self._latent_dim,self._img_shape).to(self._device) self._E = Encoder(self._latent_dim,self._img_shape).to(self._device) self._D = Discriminator(self._latent_dim,self._img_shape).to(self._device) self._G.apply(self.weights_init) self._E.apply(self.weights_init) self._D.apply(self.discriminator_weights_init) self._G_optimizer = torch.optim.Adam([{'params' : self._G.parameters()},{'params' : self._E.parameters()}], lr=self._generator_lr,betas=(0.5,0.999),weight_decay=self._weight_decay) self._D_optimizer = torch.optim.Adam(self._D.parameters(),lr=self._discriminator_lr,betas=(0.5,0.999)) self._G_scheduler = lr_scheduler.ExponentialLR(self._G_optimizer, gamma= 0.99) self._D_scheduler = lr_scheduler.ExponentialLR(self._D_optimizer, gamma= 0.99)
def build_model(self): """Build generator and discriminator.""" self.generator = Generator(z_dim=self.z_dim) print(count_parameters(self.generator)) self.discriminator = Discriminator() print(count_parameters(self.discriminator)) self.g_optimizer = optim.Adam(self.generator.parameters(), self.lr, (self.beta1, self.beta2)) self.d_optimizer = optim.Adam(self.discriminator.parameters(), self.lr*1, (self.beta1, self.beta2)) if self.epoch: g_path = os.path.join(self.model_path, 'generator-%d.pkl' % self.epoch) d_path = os.path.join(self.model_path, 'discriminator-%d.pkl' % self.epoch) g_optim_path = os.path.join(self.model_path, 'gen-optim-%d.pkl' % self.epoch) d_optim_path = os.path.join(self.model_path, 'dis-optim-%d.pkl' % self.epoch) self.generator.load_state_dict(torch.load(g_path)) self.discriminator.load_state_dict(torch.load(d_path)) self.g_optimizer.load_state_dict(torch.load(g_optim_path)) self.d_optimizer.load_state_dict(torch.load(d_optim_path)) if torch.cuda.is_available(): self.generator.cuda() self.discriminator.cuda()
def run(argv): global args global model global loader global test_loader global output_shape global visualization_mode parser = argparse.ArgumentParser() # training parser.add_argument('--train', action='store_true', default=False, help='Train a model') parser.add_argument('--train_generator', type=bool, default=True, help='Include generator training') parser.add_argument('--save_to', type=str, default='model.bin', help='Where to save the trained model') parser.add_argument('--batch_size', type=int, default=2048, help='Training batch size') parser.add_argument('--epochs', type=int, default=200, help='Number of epochs to train') parser.add_argument('--lr', type=float, default=0.001, help='Learning rate') parser.add_argument('--lr_step', type=int, default=0, help='Learning rate step size') parser.add_argument('--lr_step_gamma', type=float, default=0.5, help='Learning rate step gamma') parser.add_argument('--sigma', type=float, default=0.2, help='Noise standard deviation') parser.add_argument('--reduce_sigma_every', type=int, default=0, help='Decrease sigma value every N batches') parser.add_argument('--reduce_sigma_gamma', type=float, default=0.8, help='When decreasing sigma, sigma *= sigma_gamma') parser.add_argument("--min_sigma", type=float, default=0, help='Lower bound of sigma') parser.add_argument('--stop_training_real_dde_after', type=int, default=-1, help='Only train real_dde for this number of epochs') parser.add_argument('--train_gen_every', type=int, default=10, help='Step generator every N iterations') parser.add_argument('--visualize_every', type=int, default=50, help='Visualize model every N batches') # loading parser.add_argument('--load_from', type=str, default=None, help='Where to load a trained model') # dataset parser.add_argument('--dataset', type=str, default='ts', help='Dataset (ts, eg, cb, mnist, fashion, stacked-mnist, ncsn)') parser.add_argument('--dataset_dir', type=str, default='/tmp/', help='Dataset location for MNIST and Fashion. Defaults to /tmp/') parser.add_argument('--download', action='store_true', default=False, help='Download the dataset if it does not exist') parser.add_argument('--datasize', type=int, default=51200, help='Datasize (pseudo, for epoch computation) for 2D datasets') parser.add_argument('--dynamic', action='store_true', default=True, help='Whether 2D datasets are dynamically generated') # model parameters parser.add_argument('--dense_sigmoid', action='store_true', default=False, help='Add sigmoid to DenseNet generator') parser.add_argument('--gen_activation', type=str, default='softplus', help='Generator activation function') parser.add_argument('--dde_activation', type=str, default='softplus', help='Discriminator activation function') parser.add_argument('--n_input', type=int, default=2, help='Number of input to generator') parser.add_argument('--n_hidden', type=int, default=32, help='Number of hidden units in a layer') parser.add_argument('--n_gen_hidden', type=int, default=64, help='Number of hidden units in DenseNet generator layer') parser.add_argument('--n_layers', type=int, default=25, help='Number of layers') parser.add_argument('--n_feat_gen', type=int, default=64, help='Number of features in DCGAN generator') parser.add_argument('--n_feat_dsc', type=int, default=45, help='Number of features in DCGAN discriminator') parser.add_argument('--gen_bn', type=bool, default=True, help='Use BN in DCGAN generator') parser.add_argument('--dsc_activation', type=str, default='lrelu:0.2', help='Activation in DCGAN discriminator') parser.add_argument('--device', type=str, default='cuda:0', help='PyTorch device') # visualization parser.add_argument('--visualize', action='store_true', default=False, help='Visualize a model') parser.add_argument('--n_toy_samples', type=int, default=50000, help='Number of samples to visualize for toy datasets') parser.add_argument('--save_toy_samples_to', type=str, default=None, help='If specified, save generated toy samples to file') parser.add_argument('--vis_path', type=str, default='', help='Path prefix to save visualization figures.') parser.add_argument('--n_uci_samples', type=int, default=51200, help='Number of samples to use when computing UCI ALL') args = parser.parse_args(argv) if not argv is None else parser.parse_args() if args.dataset == 'ts': loader = twoSpiralsLoader(args.datasize, args.batch_size, dynamic=args.dynamic) args.type = 'mlp' elif args.dataset == 'eg': loader = eightGaussiansLoader(args.datasize, args.batch_size, dynamic=args.dynamic) args.type = 'mlp' elif args.dataset == 'cb': loader = checkerboardLoader(args.datasize, args.batch_size, dynamic=args.dynamic) args.type = 'mlp' elif args.dataset == 'mnist': loader = mnistLoader(args.dataset_dir, args.download, args.batch_size) args.type = 'densenet' visualization_mode = 'images' elif args.dataset == 'fashion': loader = fashionLoader(args.dataset_dir, args.download, args.batch_size) args.type = 'densenet' visualization_mode = 'images' elif args.dataset == 'stacked-mnist': loader = stackedMnistLoader(args.dataset_dir, args.download, args.batch_size) args.type = 'dcgan' visualization_mode = 'images' elif args.dataset == 'ncsn': loader = stackedMnistLoader(args.dataset_dir, args.download, args.batch_size) args.type = 'ncsn' visualization_mode = 'images' elif args.dataset.startswith('uci_'): name = args.dataset[4:] loader, test_loader = uciLoader(args.dataset_dir, name, args.batch_size) args.type = 'mlp' visualization_mode = 'numbers' else: print(f'Unknown dataset: {args.dataset}') return output_shape = list(loader.dataset[0][0].shape) print(f'Output shape: {output_shape}') output_length = loader.dataset[0][0].view(-1).shape[0] print(f'Output length: {output_length}') gen_activation = activation_from_name(args.gen_activation) dde_activation = activation_from_name(args.dde_activation) if args.type == 'mlp': model = DdeModel(MlpModule(output_length, args.n_hidden, args.n_layers, output_length, gen_activation, True, args.device), MlpModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, False, args.device), MlpModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, False, args.device)) elif args.type == 'densenet': end_activation = nn.Sigmoid if args.dense_sigmoid else None model = DdeModel(DenseNetModule(args.n_input, args.n_gen_hidden, args.n_layers, output_length, gen_activation, end_activation, args.device), DenseNetModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, None, args.device), DenseNetModule(output_length, args.n_hidden, args.n_layers, 1, dde_activation, None, args.device)) elif args.type == 'dcgan': dsc_activation = activation_from_name(args.dsc_activation) model = DdeModel(Generator(args.n_input, args.n_feat_gen, output_shape[0], args.gen_bn, args.device), Discriminator(args.n_feat_dsc, output_shape[0], dsc_activation, None, False, args.device), Discriminator(args.n_feat_dsc, output_shape[0], dsc_activation, None, False, args.device)) elif args.type == 'ncsn': model = NetModel(NCSNDde(64, gen_activation, args.device)) else: print(f'Unknown model type: {args.type}') if not args.load_from is None: model.load(args.load_from) print(f'Loaded model from {args.load_from}') if len(args.vis_path) > 0: os.makedirs(args.vis_path, exist_ok=True) if args.train: if args.type == 'ncsn': train_ncsn() else: train() if args.visualize: visualize(None, args.min_sigma)
loss_fn = nn.BCELoss() loss_rec = nn.MSELoss() controller = Segmentation(args.state_dim, args.policy_hidden_1, args.option_num) encoder = MLPEncoder(args.option_dim, 2 * args.option_dim, 1, 0, True) decoder = MLPDecoder(args.option_dim, 1, 2 * args.option_dim, 2 * args.option_dim, 2 * args.option_dim, 0, False) actor = Actor_with_option(args.state_dim, args.option_dim, args.action_dim) optimizer = optim.Adam(list(encoder.parameters()) + list(decoder.parameters()) + list(actor.parameters()) + list(controller.parameters()), lr=args.lr, betas=(0.5, 0.999)) discriminator = Discriminator(args.state_dim, args.action_dim) optimizer_discriminator = torch.optim.Adam(discriminator.parameters(), lr=1e-4, betas=(0.5, 0.999)) warm_start_optimizer = optim.Adam(list(controller.parameters()), lr=1e-2, betas=(0.5, 0.999)) if args.CUDA: policy.cuda() encoder.cuda() decoder.cuda() controller.cuda() meta_train_tasks = [0, 1, 2, 4] meta_test_tasks = [5, 6]
def model_summary(model_type, img_res, hidden_size, enc_type, dec_type, loss, batch_size, device=torch.device("cuda:1"), verbose=True): pattern = re.compile(r"Params size \(MB\):(.*)\n") pattern2 = re.compile(r"Forward/backward pass size \(MB\):(.*)\n") input_dim = 3 enc_input_size = (input_dim, img_res, img_res) dec_input_size = (hidden_size, img_res // 4, img_res // 4) pdb.set_trace() if verbose: print(f"model:{model_type}") print(f"depth:{enc_type}_{dec_type}") if model_type == "acai": model = ACAI(img_res, input_dim, hidden_size, enc_type, dec_type).to(device) elif model_type == "vqvae": model = VectorQuantizedVAE(input_dim, hidden_size, enc_type=enc_type, dec_type=dec_type).to(device) elif model_type == "vae": model = VAE(input_dim, hidden_size, enc_type=enc_type, dec_type=dec_type).to(device) encoder_summary, _ = torchsummary.summary_string(model.encoder, enc_input_size, device=device, batch_size=batch_size) decoder_summary, _ = torchsummary.summary_string(model.decoder, dec_input_size, device=device, batch_size=batch_size) if verbose: print(encoder_summary) print(decoder_summary) discriminators = {} if model_type == "acai": disc = Discriminator(input_dim, img_res, "image").to(device) disc_summary, _ = torchsummary.summary_string(disc, enc_input_size, device=device, batch_size=batch_size) disc_param_size = float(re.search(pattern, disc_summary).group(1)) disc_forward_size = float(re.search(pattern2, disc_summary).group(1)) discriminators["interp_disc"] = (disc_param_size, disc_forward_size) if loss == "gan": disc = Discriminator(input_dim, img_res, "image").to(device) disc_summary, _ = torchsummary.summary_string(disc, enc_input_size, device=device, batch_size=batch_size) disc_param_size = float(re.search(pattern, disc_summary).group(1)) disc_forward_size = float(re.search(pattern2, disc_summary).group(1)) discriminators["recons_disc"] = (disc_param_size, 2 * disc_forward_size) elif loss == "comp": disc = AnchorComparator(input_dim * 2, img_res, "image").to(device) disc_summary, _ = torchsummary.summary_string(disc, enc_input_size, device=device, batch_size=batch_size) disc_param_size = float(re.search(pattern, disc_summary).group(1)) disc_forward_size = float(re.search(pattern2, disc_summary).group(1)) discriminators["recons_disc"] = (disc_param_size, 2 * disc_forward_size) elif "comp_2" in loss: disc = ClubbedPermutationComparator(input_dim * 2, img_res, "image").to(device) disc_summary, _ = torchsummary.summary_string(disc, enc_input_size, device=device, batch_size=batch_size) disc_param_size = float(re.search(pattern, disc_summary).group(1)) disc_forward_size = float(re.search(pattern2, disc_summary).group(1)) discriminators["recons_disc"] = (disc_param_size, 2 * disc_forward_size) elif "comp_6" in loss: disc = FullPermutationComparator(input_dim * 2, img_res, "image").to(device) disc_summary, _ = torchsummary.summary_string(disc, enc_input_size, device=device, batch_size=batch_size) disc_param_size = float(re.search(pattern, disc_summary).group(1)) disc_forward_size = float(re.search(pattern2, disc_summary).group(1)) discriminators["recons_disc"] = (disc_param_size, 2 * disc_forward_size) encoder_param_size = float(re.search(pattern, encoder_summary).group(1)) encoder_forward_size = float(re.search(pattern2, encoder_summary).group(1)) decoder_param_size = float(re.search(pattern, decoder_summary).group(1)) decoder_forward_size = float(re.search(pattern2, decoder_summary).group(1)) if verbose: if "ACAI" in str(type(model)): print( f"discriminator:\n\tparams:{disc_param_size}\n\tforward:{disc_forward_size}" ) if loss == "gan": print( f"reconstruction discriminator:\n\tparams:{disc_param_size}\n\tforward:{disc_forward_size}" ) print( f"encoder:\n\tparams:{encoder_param_size}\n\tforward:{encoder_forward_size}" ) print( f"decoder:\n\tparams:{decoder_param_size}\n\tforward:{decoder_forward_size}" ) encoder = {"params": encoder_param_size, "forward": encoder_forward_size} decoder = {"params": decoder_param_size, "forward": decoder_forward_size} return encoder, decoder, discriminators
EPOCHS = 20 Z_DIM = 100 D_STEPS = 1 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') t = transforms.Compose([transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor(), lambda x: x * 2 - 1]) # Scaling to -1, 1 dataset = torchvision.datasets.CelebA('G:/Datasets', download=DOWNLOAD, transform=t) dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True) G = Generator(Z_DIM).to(device) D = Discriminator().to(device) d_loss = D_loss().to(device) g_loss = G_loss().to(device) optim_D = torch.optim.Adam(D.parameters(), lr=1e-5, betas=(0.5, 0.999)) optim_G = torch.optim.Adam(G.parameters(), lr=1e-4, betas=(0.5, 0.999)) d_count = 0 for e in range(EPOCHS): for x, _ in dataloader: for p in D.parameters(): p.requires_grad = True x = x.to(device) D.zero_grad() d_count += 1
def main(args): writer = SummaryWriter('./logs/{0}'.format(args.output_folder)) save_filename = './models/{0}'.format(args.output_folder) d_args = vars(args) pert_types = train_util.get_perturb_types(args) train_loader, valid_loader, test_loader = train_util.get_dataloaders( args, pert_types) recons_input_img = train_util.log_input_img_grid(test_loader, writer) # print(f"nn num:{len(pert_types)}") num_perturb_types = len(pert_types) input_dim = 3 model = PCIE(args.img_res, input_dim, args.hidden_size, num_perturb_types, args.enc_type, args.dec_type) interp_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) interp_disc_opt = torch.optim.Adam(interp_disc.parameters(), lr=args.disc_lr, amsgrad=True) # if torch.cuda.device_count() > 1 and args.device == "cuda": # model = torch.nn.DataParallel(model) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) # interp_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(interp_disc_opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators = {"interp_disc": [interp_disc, interp_disc_opt]} if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: if "color" in args.recons_loss: recons_disc = FullPermutationColorComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) else: recons_disc = FullPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) recons_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( recons_disc_opt, "min", patience=args.lr_patience, factor=0.5, threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] if args.prior_loss == "gan": prior_disc = Latent2ClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor) prior_disc_opt = torch.optim.Adam(prior_disc.parameters(), lr=args.disc_lr, amsgrad=True) # prior_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(prior_disc_opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators["prior_disc"] = [prior_disc, prior_disc_opt] if args.perturb_feat_gan: for perturb_type in train_util.perturb_dict: if d_args[perturb_type]: num_class = train_util.perturb_dict[perturb_type].num_class if num_class == 2: pert_disc = Latent2ClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor) pert_disc_opt = torch.optim.Adam(pert_disc.parameters(), lr=args.disc_lr, amsgrad=True) else: pert_disc = LatentMultiClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor, num_class) pert_disc_opt = torch.optim.Adam(pert_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators[f"{perturb_type}_disc"] = (pert_disc, pert_disc_opt) model.to(args.device) for disc in discriminators: discriminators[disc][0].to(args.device) # Generate the samples first once # train_util.log_recons_img_grid(recons_input_img, model, 0, args.device, writer) if args.weights == "load": start_epoch = train_util.load_state(save_filename, model, optimizer, discriminators) else: start_epoch = 0 # stop_patience = args.stop_patience best_loss = torch.tensor(np.inf) for epoch in tqdm(range(start_epoch, args.num_epochs), file=sys.stdout): # for CUDA OOM error, prevents running dependency job on slurm which is meant to run on timeout try: train(epoch, train_loader, model, optimizer, args, writer, discriminators) # pass except RuntimeError as err: print("".join( traceback.TracebackException.from_exception(err).format()), file=sys.stderr) print("*******", file=sys.stderr) print(err, file=sys.stderr) exit(0) print("out of train") # comp = subprocess.run("nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits", text=True, stdout=subprocess.PIPE) # print(comp.stdout, file=sys.stderr) val_loss_dict, _ = ae_test(get_losses, model, valid_loader, args, discriminators) train_util.log_losses("val", val_loss_dict, epoch + 1, writer) # print("logg loss") # train_util.log_latent_metrics("val", z, epoch+1, writer) # print("log metric") train_util.save_recons_img_grid("test", recons_input_img, model, epoch + 1, args) # print("log recons") train_util.save_interp_img_grid("test", recons_input_img, model, epoch + 1, args) # print("log interp") train_util.save_state(model, optimizer, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, save_filename)
def main(args): input_dim = 3 model = VAE(input_dim, args.hidden_size, args.enc_type, args.dec_type) opt = torch.optim.Adam(model.parameters(), lr=args.lr, amsgrad=True) # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-6) # ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators = {} if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: recons_disc = FullPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] if torch.cuda.device_count() > 1: model = train_util.ae_data_parallel(model) for disc in discriminators: discriminators[disc][0] = torch.nn.DataParallel( discriminators[disc][0]) model.to(args.device) for disc in discriminators: discriminators[disc][0].to(args.device) print("model built", file=sys.stderr) #print("model created") train_loader, val_loader, test_loader = train_util.get_dataloaders(args) print("loaders acquired", file=sys.stderr) #print("loaders acquired") model_name = f"vae_{args.recons_loss}" if args.output_folder is None: args.output_folder = os.path.join( model_name, args.dataset, f"depth_{args.enc_type}_{args.dec_type}_hs_{args.img_res}_{args.hidden_size}" ) log_save_path = os.path.join("./logs", args.output_folder) model_save_path = os.path.join("./models", args.output_folder) if not os.path.exists(log_save_path): os.makedirs(log_save_path) print(f"log:{log_save_path}", file=sys.stderr) sys.stderr.flush() if not os.path.exists(model_save_path): os.makedirs(model_save_path) writer = SummaryWriter(log_save_path) print(f"train loader length:{len(train_loader)}", file=sys.stderr) best_loss = torch.tensor(np.inf) if args.weights == "load": start_epoch = train_util.load_state(model_save_path, model, opt, discriminators) else: start_epoch = 0 recons_input_img = train_util.log_input_img_grid(test_loader, writer) train_util.log_recons_img_grid(recons_input_img, model, 0, args.device, writer) stop_patience = args.stop_patience for epoch in range(start_epoch, args.num_epochs): try: train(model, train_loader, opt, epoch, writer, args, discriminators) except RuntimeError as err: print("".join( traceback.TracebackException.from_exception(err).format()), file=sys.stderr) print("*******", file=sys.stderr) print(err, file=sys.stderr) exit(0) val_loss_dict, z = train_util.test(get_losses, model, val_loader, args, discriminators) print(f"epoch loss:{val_loss_dict['recons_loss'].item()}") train_util.save_recons_img_grid("test", recons_input_img, model, epoch + 1, args) train_util.save_interp_img_grid("test", recons_input_img, model, epoch + 1, args) train_util.log_losses("val", val_loss_dict, epoch + 1, writer) train_util.log_latent_metrics("val", z, epoch + 1, writer) train_util.save_state(model, opt, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, model_save_path)
def main(args): train_loader, val_loader, test_loader = train_util.get_dataloaders(args) input_dim = 3 model = VAE(input_dim, args.hidden_size, args.enc_type, args.dec_type) opt = torch.optim.Adam(model.parameters(), lr=LR, amsgrad=True) discriminators = {} if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim*2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator(input_dim*2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: recons_disc = FullPermutationComparator(input_dim*2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] if torch.cuda.device_count() > 1: model = train_util.ae_data_parallel(model) for disc in discriminators: discriminators[disc][0] = torch.nn.DataParallel(discriminators[disc][0]) model.to(args.device) model_name = f"vae_{args.recons_loss}" if args.output_folder is None: args.output_folder = os.path.join(model_name, args.dataset, f"depth_{args.enc_type}_{args.dec_type}_hs_{args.img_res}_{args.hidden_size}") log_save_path = os.path.join("./logs", args.output_folder) model_save_path = os.path.join("./models", args.output_folder) if not os.path.exists(log_save_path): os.makedirs(log_save_path) print(f"log:{log_save_path}", file=sys.stderr) sys.stderr.flush() if not os.path.exists(model_save_path): os.makedirs(model_save_path) writer = SummaryWriter(log_save_path) print(f"train loader length:{len(train_loader)}", file=sys.stderr) best_loss = torch.tensor(np.inf) if args.weights == "load": start_epoch = train_util.load_state(model_save_path, model, opt, discriminators) else: start_epoch = 0 recons_input_img = train_util.log_input_img_grid(test_loader, writer) train_util.save_recons_img_grid("val", recons_input_img, model, 0, args) for epoch in range(1, args.num_epochs): print("Epoch {}:".format(epoch)) train(model, opt, train_loader) curr_loss = val(model, val_loader) # val_loss_dict, z = train_util.test(get_losses, model, val_loader, args, discriminators) print(f"epoch val loss:{curr_loss}", file=sys.stderr) sys.stderr.flush() train_util.save_recons_img_grid("val", recons_input_img, model, epoch+1, args) train_util.save_interp_img_grid("val", recons_input_img, model, epoch+1, args)
def __init__(self, n_input: int, n_batch: int = 0, n_labels: int = 0, n_hidden: int = 256, n_latent: int = 18, n_layers: int = 1, dropout_rate: float = 0., dispersion: str = "gene", log_variational: bool = False, reconstruction_loss: str = "alpha-gan", n_centroids=12, X=None, gan_loss='gan', reconst_ratio=0.01, use_cuda: bool = True, n_batch_gan: int = 0): super().__init__() self.dispersion = dispersion self.n_latent = n_latent self.n_hidden = n_hidden self.n_layers = n_layers self.log_variational = log_variational self.reconstruction_loss = reconstruction_loss # Automatically deactivate if useless self.n_batch = n_batch self.n_labels = n_labels self.n_centroids = n_centroids self.dropout_rate = dropout_rate self.X = X self.gan_loss = gan_loss self.reconst_ratio = reconst_ratio self.use_cuda = use_cuda self.n_batch_gan = n_batch_gan if self.dispersion == "gene": self.px_r = torch.nn.Parameter(torch.randn(n_input)) elif self.dispersion == "gene-batch": # batch is different times of exp self.px_r = torch.nn.Parameter(torch.randn(n_input, n_batch)) elif self.dispersion == "gene-label": self.px_r = torch.nn.Parameter(torch.randn(n_input, n_labels)) elif self.dispersion == "gene-cell": pass else: raise ValueError("dispersion must be one of ['gene', 'gene-batch'," " 'gene-label', 'gene-cell'], but input was " "{}.format(self.dispersion)") self.pi = nn.Parameter(torch.ones(n_centroids) / n_centroids) # pc self.mu_c = nn.Parameter(torch.zeros(n_latent, n_centroids)) # mu self.var_c = nn.Parameter(torch.ones(n_latent, n_centroids)) # sigma^2 # z encoder goes from the n_input-dimensional data to an n_latent-d # latent space representation self.z_encoder = Encoder(n_input, n_latent, n_layers=n_layers, n_hidden=n_hidden, dropout_rate=dropout_rate) # l encoder goes from n_input-dimensional data to 1-d library size self.l_encoder = Encoder(n_input, 1, n_layers=1, n_hidden=n_hidden, dropout_rate=dropout_rate) # discriminator to distinguish the generated/fake/real and batches self.discriminator = Discriminator(n_input, n_hidden, self.n_batch_gan, gan_loss) # decoder goes from n_latent-dimensional space to n_input-d data self.decoder = DecoderSCVI( n_latent, n_input, n_cat_list=[n_batch], n_layers=n_layers, n_hidden=n_hidden, )
def main(args): writer = SummaryWriter('./logs/{0}'.format(args.output_folder)) save_filename = './models/{0}'.format(args.output_folder) train_loader, valid_loader, test_loader = train_util.get_dataloaders(args) num_channels = 3 model = VectorQuantizedVAE(num_channels, args.hidden_size, args.k, args.enc_type, args.dec_type) model.to(args.device) # Fixed images for Tensorboard recons_input_img = train_util.log_input_img_grid(test_loader, writer) train_util.log_recons_img_grid(recons_input_img, model, 0, args.device, writer) discriminators = {} input_dim = 3 if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: recons_disc = FullPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) recons_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( recons_disc_opt, "min", patience=args.lr_patience, factor=0.5, threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, "min", patience=args.lr_patience, factor=0.5, threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) if torch.cuda.device_count() > 1: model = train_util.ae_data_parallel(model) for disc in discriminators: discriminators[disc][0] = torch.nn.DataParallel( discriminators[disc][0]) model.to(args.device) for disc in discriminators: discriminators[disc][0].to(args.device) # Generate the samples first once recons_input_img = train_util.log_input_img_grid(test_loader, writer) train_util.log_recons_img_grid(recons_input_img, model, 0, args.device, writer) if args.weights == "load": start_epoch = train_util.load_state(save_filename, model, optimizer, discriminators) else: start_epoch = 0 stop_patience = args.stop_patience best_loss = torch.tensor(np.inf) for epoch in tqdm(range(start_epoch, args.num_epochs), file=sys.stdout): try: train(epoch, train_loader, model, optimizer, args, writer, discriminators) except RuntimeError as err: print("".join( traceback.TracebackException.from_exception(err).format()), file=sys.stderr) print("*******") print(err, file=sys.stderr) print(f"batch_size:{args.batch_size}", file=sys.stderr) exit(0) val_loss_dict, z = train_util.test(get_losses, model, valid_loader, args, discriminators) train_util.log_recons_img_grid(recons_input_img, model, epoch + 1, args.device, writer) train_util.log_interp_img_grid(recons_input_img, model, epoch + 1, args.device, writer) train_util.log_losses("val", val_loss_dict, epoch + 1, writer) train_util.log_latent_metrics("val", z, epoch + 1, writer) train_util.save_state(model, optimizer, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, save_filename) # early stop check # if val_loss_dict["recons_loss"] - best_loss < args.threshold: # stop_patience -= 1 # else: # stop_patience = args.stop_patience # if stop_patience == 0: # print("training early stopped!") # break ae_lr_scheduler.step(val_loss_dict["recons_loss"]) if args.recons_loss != "mse": recons_disc_lr_scheduler.step(val_loss_dict["recons_disc_loss"])
train=True, download=True, transform=transforms.ToTensor()), batch_size=64, shuffle=False, **kwargs) test_loader = torch.utils.data.DataLoader(datasets.MNIST( '../data/fashion_mnist/', train=False, transform=transforms.ToTensor()), batch_size=32, shuffle=False, **kwargs) test_data = list(test_loader) netG = Generator().cuda() netD = Discriminator().cuda() optimizerG = torch.optim.Adam(netG.parameters(), lr=1e-4, betas=(0.5, 0.9)) optimizerD = torch.optim.Adam(netD.parameters(), lr=1e-4, betas=(0.5, 0.9)) one = torch.cuda.FloatTensor([1]) mone = one * -1 def train(epoch): train_loss = [] for batch_idx, (data, _) in enumerate(train_loader): start_time = time.time() if data.size(0) != 64: continue
def train(rank: int, cfg: DictConfig): print(OmegaConf.to_yaml(cfg)) if cfg.train.n_gpu > 1: init_process_group(backend=cfg.train.dist_config['dist_backend'], init_method=cfg.train.dist_config['dist_url'], world_size=cfg.train.dist_config['world_size'] * cfg.train.n_gpu, rank=rank) device = torch.device( 'cuda:{:d}'.format(rank) if torch.cuda.is_available() else 'cpu') generator = Generator(sum(cfg.model.feature_dims), *cfg.model.cond_dims, **cfg.model.generator).to(device) discriminator = Discriminator(**cfg.model.discriminator).to(device) if rank == 0: print(generator) os.makedirs(cfg.train.ckpt_dir, exist_ok=True) print("checkpoints directory : ", cfg.train.ckpt_dir) if os.path.isdir(cfg.train.ckpt_dir): cp_g = scan_checkpoint(cfg.train.ckpt_dir, 'g_') cp_do = scan_checkpoint(cfg.train.ckpt_dir, 'd_') steps = 1 if cp_g is None or cp_do is None: state_dict_do = None last_epoch = -1 else: state_dict_g = load_checkpoint(cp_g, device) state_dict_do = load_checkpoint(cp_do, device) generator.load_state_dict(state_dict_g['generator']) discriminator.load_state_dict(state_dict_do['discriminator']) steps = state_dict_do['steps'] + 1 last_epoch = state_dict_do['epoch'] if cfg.train.n_gpu > 1: generator = DistributedDataParallel(generator, device_ids=[rank]).to(device) discriminator = DistributedDataParallel(discriminator, device_ids=[rank]).to(device) optim_g = RAdam(generator.parameters(), cfg.opt.lr, betas=cfg.opt.betas) optim_d = RAdam(discriminator.parameters(), cfg.opt.lr, betas=cfg.opt.betas) if state_dict_do is not None: optim_g.load_state_dict(state_dict_do['optim_g']) optim_d.load_state_dict(state_dict_do['optim_d']) scheduler_g = torch.optim.lr_scheduler.ExponentialLR( optim_g, gamma=cfg.opt.lr_decay, last_epoch=last_epoch) scheduler_d = torch.optim.lr_scheduler.ExponentialLR( optim_d, gamma=cfg.opt.lr_decay, last_epoch=last_epoch) train_filelist = load_dataset_filelist(cfg.dataset.train_list) trainset = FeatureDataset(cfg.dataset, train_filelist, cfg.data) train_sampler = DistributedSampler( trainset) if cfg.train.n_gpu > 1 else None train_loader = DataLoader(trainset, batch_size=cfg.train.batch_size, num_workers=cfg.train.num_workers, shuffle=True, sampler=train_sampler, pin_memory=True, drop_last=True) if rank == 0: val_filelist = load_dataset_filelist(cfg.dataset.test_list) valset = FeatureDataset(cfg.dataset, val_filelist, cfg.data, segmented=False) val_loader = DataLoader(valset, batch_size=1, num_workers=cfg.train.num_workers, shuffle=False, sampler=train_sampler, pin_memory=True) sw = SummaryWriter(os.path.join(cfg.train.ckpt_dir, 'logs')) generator.train() discriminator.train() for epoch in range(max(0, last_epoch), cfg.train.epochs): if rank == 0: start = time.time() print("Epoch: {}".format(epoch + 1)) if cfg.train.n_gpu > 1: train_sampler.set_epoch(epoch) for y, x_noised_features, x_noised_cond in train_loader: if rank == 0: start_b = time.time() y = y.to(device, non_blocking=True) x_noised_features = x_noised_features.transpose(1, 2).to( device, non_blocking=True) x_noised_cond = x_noised_cond.to(device, non_blocking=True) z1 = torch.randn(cfg.train.batch_size, cfg.model.cond_dims[1], device=device) z2 = torch.randn(cfg.train.batch_size, cfg.model.cond_dims[1], device=device) y_hat1 = generator(x_noised_features, x_noised_cond, z=z1) y_hat2 = generator(x_noised_features, x_noised_cond, z=z2) # Discriminator real_scores, fake_scores = discriminator(y), discriminator( y_hat1.detach()) d_loss = discriminator_loss(real_scores, fake_scores) optim_d.zero_grad() d_loss.backward(retain_graph=True) optim_d.step() # Generator g_stft_loss = criterion(y, y_hat1) + criterion( y, y_hat2) - criterion(y_hat1, y_hat2) g_adv_loss = adversarial_loss(fake_scores) g_loss = g_adv_loss + g_stft_loss optim_g.zero_grad() g_loss.backward() optim_g.step() if rank == 0: # STDOUT logging if steps % cfg.train.stdout_interval == 0: with torch.no_grad(): print( 'Steps : {:d}, Gen Loss Total : {:4.3f}, STFT Error : {:4.3f}, s/b : {:4.3f}' .format(steps, g_loss, g_stft_loss, time.time() - start_b)) # checkpointing if steps % cfg.train.checkpoint_interval == 0: ckpt_dir = "{}/g_{:08d}".format(cfg.train.ckpt_dir, steps) save_checkpoint( ckpt_dir, { 'generator': (generator.module if cfg.train.n_gpu > 1 else generator).state_dict() }) ckpt_dir = "{}/do_{:08d}".format(cfg.train.ckpt_dir, steps) save_checkpoint( ckpt_dir, { 'discriminator': (discriminator.module if cfg.train.n_gpu > 1 else discriminator).state_dict(), 'optim_g': optim_g.state_dict(), 'optim_d': optim_d.state_dict(), 'steps': steps, 'epoch': epoch }) # Tensorboard summary logging if steps % cfg.train.summary_interval == 0: sw.add_scalar("training/gen_loss_total", g_loss, steps) sw.add_scalar("training/gen_stft_error", g_stft_loss, steps) # Validation if steps % cfg.train.validation_interval == 0: generator.eval() torch.cuda.empty_cache() val_err_tot = 0 with torch.no_grad(): for j, (y, x_noised_features, x_noised_cond) in enumerate(val_loader): y_hat = generator( x_noised_features.transpose(1, 2).to(device), x_noised_cond.to(device)) val_err_tot += criterion(y, y_hat).item() if j <= 4: # sw.add_audio('noised/y_noised_{}'.format(j), y_noised[0], steps, cfg.data.target_sample_rate) sw.add_audio('generated/y_hat_{}'.format(j), y_hat[0], steps, cfg.data.sample_rate) sw.add_audio('gt/y_{}'.format(j), y[0], steps, cfg.data.sample_rate) val_err = val_err_tot / (j + 1) sw.add_scalar("validation/stft_error", val_err, steps) generator.train() steps += 1 scheduler_g.step() scheduler_d.step() if rank == 0: print('Time taken for epoch {} is {} sec\n'.format( epoch + 1, int(time.time() - start)))
TF2_test_data_source, mfcc_only=False, norm_calc=False) test_dataset.input_meanstd = train_dataset.input_meanstd test_dataset.output_meanstd = train_dataset.output_meanstd train_dataset_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_dataset_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) generator_A2B = Generator(24).to("cuda") generator_B2A = Generator(24).to("cuda") discriminator_A = Discriminator(1).to("cuda") discriminator_B = Discriminator(1).to("cuda") generator_params = [generator_A2B.parameters(), generator_B2A.parameters()] generator_optimizer = torch.optim.Adam(itertools.chain(*generator_params), lr=generator_lr) discriminator_params = [ discriminator_A.parameters(), discriminator_B.parameters() ] discriminator_optimizer = torch.optim.Adam( itertools.chain(*discriminator_params), lr=discriminator_lr) for epoch in range(num_epochs): print("Epoch ", epoch) for i, sample in enumerate(train_dataset_loader):
opt = parser.parse_args() # Write learning info writer = SummaryWriter('runs/2d_gaussain_mixture') np.random.seed(opt.random_seed) torch.manual_seed(opt.random_seed) cuda = True if torch.cuda.is_available() and opt.cuda else False Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() and opt.cuda else torch.FloatTensor # Loss weight for gradient penalty lambda_gp = opt.lambda_gp # Initialize generator and discriminator generator = Generator(opt.latent_dim, 2) discriminator = Discriminator(2) gaussian = GaussianMixture(opt.grid_side_size) train_data = gaussian.sample(opt.batch_size * opt.target_factor if opt.learning_type == 'smallgan' else opt.batch_size) if cuda: generator.cuda() discriminator.cuda() prior_batch_size = opt.batch_size * opt.prior_factor if opt.learning_type == 'smallgan' else opt.batch_size # Valid dataset valid_samples = np.random.uniform(-1, 1, (opt.valid_size, opt.latent_dim)) valid_dataloader = torch.utils.data.DataLoader( valid_samples, batch_size=opt.valid_batch_size, shuffle=False,
def discriminator(cfg): discriminator = Discriminator(**cfg.model.discriminator) return discriminator
def main(args): writer = SummaryWriter('./logs/{0}'.format(args.output_folder)) save_filename = './models/{0}'.format(args.output_folder) train_loader, val_loader, test_loader = train_util.get_dataloaders(args) recons_input_img = train_util.log_input_img_grid(test_loader, writer) input_dim = 3 model = ACAI(args.img_res, input_dim, args.hidden_size, args.enc_type, args.dec_type).to(args.device) disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) disc_opt = torch.optim.Adam(disc.parameters(), lr=args.disc_lr, amsgrad=True) # if torch.cuda.device_count() > 1 and args.device == "cuda": # model = torch.nn.DataParallel(model) opt = torch.optim.Adam(model.parameters(), lr=args.lr) # ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) # interp_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(disc_opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators = {"interp_disc": [disc, disc_opt]} if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: recons_disc = FullPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] # Generate the samples first once train_util.save_recons_img_grid("test", recons_input_img, model, 0, args) if args.weights == "load": start_epoch = train_util.load_state(save_filename, model, opt, discriminators) else: start_epoch = 0 best_loss = torch.tensor(np.inf) for epoch in range(args.num_epochs): print("Epoch {}:".format(epoch)) train(model, opt, train_loader, args, discriminators, writer) # curr_loss = val(model, val_loader) # print(f"epoch val loss:{curr_loss}") val_loss_dict, z = train_util.test(get_losses, model, val_loader, args, discriminators) train_util.log_losses("val", val_loss_dict, epoch + 1, writer) train_util.log_latent_metrics("val", z, epoch + 1, writer) # train_util.log_recons_img_grid(recons_input_img, model, epoch+1, args.device, writer) # train_util.log_interp_img_grid(recons_input_img, model, epoch+1, args.device, writer) train_util.save_recons_img_grid("val", recons_input_img, model, epoch + 1, args) train_util.save_interp_img_grid("val", recons_input_img, model, epoch + 1, args) train_util.save_state(model, opt, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, save_filename)
def test_gan(self): # models settings lambda_gp = 0.1 Tensor = torch.FloatTensor generator = Generator(2, 2) discriminator = Discriminator(2) optimizer_G = torch.optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999)) criterion = torch.nn.BCEWithLogitsLoss() optimizer_D.zero_grad() # run models for test batch real_data = Tensor(np.random.normal(0, 1, (64, 2))) z = Tensor(np.random.normal(0, 1, (64, 2))) real_target = Tensor(real_data.size(0), 1).fill_(1.0) fake_target = Tensor(real_data.size(0), 1).fill_(0.0) g_before = deepcopy(generator) d_before = deepcopy(discriminator) # Generate a batch of images fake_data = generator(z) # Real images real_validity = discriminator(real_data) # Fake images fake_validity = discriminator(fake_data) # Gradient penalty gradient_penalty = compute_gradient_penalty(discriminator, real_data.data, fake_data.data, Tensor) # Discriminator loss d_loss = criterion(real_validity, real_target) \ + criterion(fake_validity, fake_target) \ + lambda_gp * gradient_penalty d_loss.backward() optimizer_D.step() # Assert that D changed and G not changed g_changed = [ torch.equal(after, before) for after, before in zip( generator.parameters(), g_before.parameters()) ] self.assertTrue(all(g_changed)) d_changed = [ torch.equal(after, before) for after, before in zip( discriminator.parameters(), d_before.parameters()) ] self.assertFalse(all(d_changed)) optimizer_G.zero_grad() # Train on fake samples g_before = deepcopy(generator) d_before = deepcopy(discriminator) fake_data = generator(z) fake_validity = discriminator(fake_data) g_loss = criterion(fake_validity, real_target) g_loss.backward() optimizer_G.step() # Assert that G changed and D not changed g_changed = [ torch.equal(after, before) for after, before in zip( generator.parameters(), g_before.parameters()) ] self.assertFalse(all(g_changed)) d_changed = [ torch.equal(after, before) for after, before in zip( discriminator.parameters(), d_before.parameters()) ] self.assertTrue(all(d_changed))
def val_test(args): writer = SummaryWriter('./logs/{0}'.format(args.output_folder)) save_filename = './models/{0}'.format(args.output_folder) d_args = vars(args) num_perturb_types = 0 perturb_types = [] for perturb_type in train_util.perturb_dict: if d_args[perturb_type]: num_perturb_types += 1 perturb_types.append(perturb_type) train_loader, valid_loader, test_loader = train_util.get_dataloaders( args, perturb_types) recons_input_img = train_util.log_input_img_grid(test_loader, writer) input_dim = 3 model = PCIE(args.img_res, input_dim, args.hidden_size, num_perturb_types, args.enc_type, args.dec_type) interp_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) interp_disc_opt = torch.optim.Adam(interp_disc.parameters(), lr=args.disc_lr, amsgrad=True) # if torch.cuda.device_count() > 1 and args.device == "cuda": # model = torch.nn.DataParallel(model) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # ae_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) # interp_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(interp_disc_opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators = {"interp_disc": [interp_disc, interp_disc_opt]} if args.recons_loss != "mse": if args.recons_loss == "gan": recons_disc = Discriminator(input_dim, args.img_res, args.input_type).to(args.device) elif args.recons_loss == "comp": recons_disc = AnchorComparator(input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_2" in args.recons_loss: recons_disc = ClubbedPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) elif "comp_6" in args.recons_loss: recons_disc = FullPermutationComparator( input_dim * 2, args.img_res, args.input_type).to(args.device) recons_disc_opt = torch.optim.Adam(recons_disc.parameters(), lr=args.disc_lr, amsgrad=True) recons_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( recons_disc_opt, "min", patience=args.lr_patience, factor=0.5, threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators["recons_disc"] = [recons_disc, recons_disc_opt] if args.prior_loss == "gan": prior_disc = Latent2ClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor) prior_disc_opt = torch.optim.Adam(prior_disc.parameters(), lr=args.disc_lr, amsgrad=True) # prior_disc_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(prior_disc_opt, "min", patience=args.lr_patience, factor=0.5, # threshold=args.threshold, threshold_mode="abs", min_lr=1e-7) discriminators["prior_disc"] = [prior_disc, prior_disc_opt] print("pertrub gans") if args.perturb_feat_gan: for perturb_type in train_util.perturb_dict: if d_args[perturb_type]: num_classes = train_util.perturb_dict[perturb_type].num_class pdb.set_trace() if num_classes == 2: print(f"perturb:{d_args[perturb_type]}\ttype: two ") pert_disc = Latent2ClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor) pert_disc_opt = torch.optim.Adam(pert_disc.parameters(), lr=args.disc_lr, amsgrad=True) else: print(f"perturb:{d_args[perturb_type]}\ttype: multi ") pert_disc = LatentMultiClassDiscriminator( args.hidden_size, args.img_res // args.scale_factor, num_classes) pert_disc_opt = torch.optim.Adam(pert_disc.parameters(), lr=args.disc_lr, amsgrad=True) discriminators[f"{perturb_type}_disc"] = (pert_disc, pert_disc_opt) print("perrturb gans set") model.to(args.device) for disc in discriminators: discriminators[disc][0].to(args.device) # Generate the samples first once # train_util.log_recons_img_grid(recons_input_img, model, 0, args.device, writer) if args.weights == "load": start_epoch = train_util.load_state(save_filename, model, optimizer, discriminators) else: start_epoch = 0 # stop_patience = args.stop_patience best_loss = torch.tensor(np.inf) for epoch in tqdm(range(start_epoch, 4), file=sys.stdout): val_loss_dict = train_util.test(get_losses, model, valid_loader, args, discriminators, True) if args.weights == "init" and epoch == 1: epoch += 1 break train_util.log_losses("val", val_loss_dict, epoch + 1, writer) train_util.log_recons_img_grid(recons_input_img, model, epoch, args.device, writer) train_util.log_interp_img_grid(recons_input_img, model, epoch + 1, args.device, writer) train_util.save_state(model, optimizer, discriminators, val_loss_dict["recons_loss"], best_loss, args.recons_loss, epoch, save_filename) print(val_loss_dict)