def test_trainer_acgan(self): network_params = { "generator": { "name": ACGANGenerator, "args": {"num_classes": 10, "out_channels": 1, "step_channels": 4}, "optimizer": { "name": Adam, "args": {"lr": 0.0002, "betas": (0.5, 0.999)}, }, }, "discriminator": { "name": ACGANDiscriminator, "args": {"num_classes": 10, "in_channels": 1, "step_channels": 4}, "optimizer": { "name": Adam, "args": {"lr": 0.0002, "betas": (0.5, 0.999)}, }, }, } losses_list = [ MinimaxGeneratorLoss(), MinimaxDiscriminatorLoss(), AuxiliaryClassifierGeneratorLoss(), AuxiliaryClassifierDiscriminatorLoss(), ] trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), ) trainer(mnist_dataloader())
def test_trainer_dcgan(self): network_params = { "generator": {"name": DCGANGenerator, "args": {"out_channels": 1, "step_channels": 4}, "optimizer": {"name": Adam, "args": {"lr": 0.0002, "betas": (0.5, 0.999)}}}, "discriminator": {"name": DCGANDiscriminator, "args": {"in_channels": 1, "step_channels": 4}, "optimizer": {"name": Adam, "args": {"lr": 0.0002, "betas": (0.5, 0.999)}}} } losses_list = [MinimaxGeneratorLoss(), MinimaxDiscriminatorLoss()] trainer = Trainer(network_params, losses_list, sample_size=1, epochs=1, device=torch.device('cpu')) trainer(mnist_dataloader())
def test_trainer_infogan(self): network_params = { "generator": { "name": InfoGANGenerator, "args": { "out_channels": 1, "step_channels": 4, "dim_dis": 128, "dim_cont": 128 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, "discriminator": { "name": InfoGANDiscriminator, "args": { "in_channels": 1, "step_channels": 4, "dim_dis": 128, "dim_cont": 128 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, } losses_list = [MutualInformationPenalty()] dis_code = torch.zeros(128, 128) dis_code[:, 0] = 1 trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), dis_code=dis_code, cont_code=torch.rand(128, 128), ) trainer(mnist_dataloader())
def test_trainer_cgan(self): network_params = { "generator": { "name": ConditionalGANGenerator, "args": { "num_classes": 10, "out_channels": 1, "step_channels": 4, }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, "discriminator": { "name": ConditionalGANDiscriminator, "args": { "num_classes": 10, "in_channels": 1, "step_channels": 4, }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, } losses_list = [ MinimaxGeneratorLoss(), MinimaxDiscriminatorLoss(), WassersteinGeneratorLoss(), WassersteinDiscriminatorLoss(), WassersteinGradientPenalty(), ] trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), ) trainer(mnist_dataloader())
def test_trainer_dcgan(self): network_params = { "generator": { "name": DCGANGenerator, "args": { "out_channels": 1, "step_channels": 4 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, "discriminator": { "name": DCGANDiscriminator, "args": { "in_channels": 1, "step_channels": 4 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, } losses_list = [ BoundaryEquilibriumGeneratorLoss(), HistoricalAverageGeneratorLoss(), FeatureMatchingGeneratorLoss(), BoundaryEquilibriumDiscriminatorLoss(), HistoricalAverageDiscriminatorLoss(), ] trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), ) trainer(mnist_dataloader())
def test_trainer_autoencoding_models(self): network_params = { "generator": { "name": AutoEncodingGenerator, "args": { "out_channels": 1, "step_channels": 4 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, "discriminator": { "name": AutoEncodingDiscriminator, "args": { "in_channels": 1, "step_channels": 4, "energy": False, }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, } losses_list = [ BoundaryEquilibriumGeneratorLoss(), BoundaryEquilibriumDiscriminatorLoss(), DraganGradientPenalty(), ] trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), ) trainer(mnist_dataloader())
def test_trainer_autoencoding_eb_models(self): network_params = { "generator": { "name": AutoEncodingGenerator, "args": { "out_channels": 1, "step_channels": 4 }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, "discriminator": { "name": AutoEncodingDiscriminator, "args": { "in_channels": 1, "step_channels": 4, "embeddings": True, }, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999) }, }, }, } losses_list = [ EnergyBasedGeneratorLoss(), EnergyBasedPullingAwayTerm(), EnergyBasedDiscriminatorLoss(), ] trainer = Trainer( network_params, losses_list, sample_size=1, epochs=1, device=torch.device("cpu"), ) trainer(mnist_dataloader())
batch_size = 32 if torch.cuda.is_available(): device = torch.device('cuda') epochs = 2000 else: device = torch.device("cpu") epochs = 5 print('Device:', device) print('Epochs:', epochs) print('Dataset Total Size:', dataset_size) print('Dataset For Validation Remaining:', split) print('Batch Size:', batch_size) dataloader = data.DataLoader(dataset, batch_size=batch_size, sampler=train_sampler) lsgan_losses = [ LeastSquaresGeneratorLoss(), LeastSquaresDiscriminatorLoss() ] trainer = Trainer(dcgan_network, lsgan_losses, sample_size=64, epochs=epochs, device=device) trainer(dataloader)
epochs=args.epochs, sample_size=args.sample_size, checkpoints=args.checkpoint, retain_checkpoints=1, recon=args.reconstructions, ) else: if args.cpu == 1: device = torch.device("cpu") else: device = torch.device("cuda:0") trainer = Trainer( network_config, losses_list, device=device, epochs=args.epochs, sample_size=args.sample_size, checkpoints=args.checkpoint, retain_checkpoints=1, recon=args.reconstructions, ) # Transforms to get Binarized MNIST dataset = dsets.MNIST(root=args.data_dir, train=True, transform=transforms.Compose([ transforms.Resize((32, 32)), transforms.Lambda(lambda x: x.convert('1')), transforms.ToTensor() ]), download=True)
epochs = 20 else: device = torch.device("cpu") epochs = 20 print("Device: {}".format(device)) print("Epochs: {}".format(epochs)) losses_list = [ WassersteinGeneratorLoss(), WassersteinDiscriminatorLoss(clip=(-0.01, 0.01)), ] trainer = Trainer( network_params, losses_list, sample_size=64, epochs=epochs, device=device ) trainer(dataloader) trainer.complete() # Grab a batch of real images from the dataloader real_batch = next(iter(dataloader)) # Plot the real images plt.figure(figsize=(10, 10)) plt.subplot(1, 2, 1) plt.axis("off") plt.title("Real Images")
if torch.cuda.is_available(): device = torch.device("cuda:0") # Use deterministic cudnn algorithms torch.backends.cudnn.deterministic = True epochs = 10 else: device = torch.device("cpu") epochs = 10 print("Device: {}".format(device)) print("Epochs: {}".format(epochs)) trainer = Trainer(network, losses, sample_size=64, epochs=epochs, device=device) trainer(loader) # Grab a batch of real images from the dataloader real_batch = next(iter(loader)) # Plot the real images plt.figure(figsize=(10, 10)) plt.subplot(1, 2, 1) plt.axis("off") plt.title("Real Images") plt.imshow( np.transpose( vutils.make_grid(real_batch[0].to(device)[:64],
}, "optimizer": {"name": Adam, "args": {"lr": 0.0002, "betas": (0.5, 0.999)}}, }, "discriminator": { "name": DCGANDiscriminator, "args": {"in_channels": 3, "step_channels": 16}, "optimizer": { "name": Adam, "args": { "lr": 0.0002, "betas": (0.5, 0.999), }, }, }, } loss = [MinimaxGeneratorLoss(), MinimaxDiscriminatorLoss()] trainer = Trainer( dcgan, loss, sample_size=64, epochs=100, checkpoints=_C.OUTPUT.CHECKPOINT_DIR, recon=_C.OUTPUT.PREDICTION_DIR, log_dir=_C.OUTPUT.LOG_DIR, device=device, ) trainer(cifar10())
print('Checkpoint Prefix:', checkpoint_prefix) images_dir = Path('./sample_images') if not images_dir.is_dir(): images_dir.mkdir() print('Samples image dir:', images_dir) log_dir = Path('./log/ACGAN') if not log_dir.is_dir(): log_dir.mkdir() print('Log dir:', log_dir) dataloader = data.DataLoader(dataset, batch_size=batch_size, sampler=train_sampler) lsgan_losses = [AuxiliaryClassifierGeneratorLoss(), AuxiliaryClassifierDiscriminatorLoss()] trainer = Trainer( dcgan_network, lsgan_losses, sample_size=64, epochs=epochs, device=device, checkpoints=checkpoint_prefix, recon=images_dir, log_dir=log_dir, ) print('Continueing from epoch 2000') trainer.load_model('model/ACGAN_goodfiles4.model') trainer(dataloader)