def main(args):
"""
Main function for the script
:param args: parsed command line arguments
:return: None
"""
from attn_gan_pytorch.Utils import get_layer
from attn_gan_pytorch.ConfigManagement import get_config
from attn_gan_pytorch.Networks import Generator, Discriminator, GAN
from data_processing.DataLoader import FlatDirectoryImageDataset, \
get_transform, get_data_loader
from attn_gan_pytorch.Losses import RelativisticAverageHingeGAN
# create a data source:
celeba_dataset = FlatDirectoryImageDataset(args.images_dir,
transform=get_transform((64, 64)))
data = get_data_loader(celeba_dataset, args.batch_size, args.num_workers)
# create generator object:
gen_conf = get_config(args.generator_config)
gen_conf = list(map(get_layer, gen_conf.architecture))
generator = Generator(gen_conf, args.latent_size)
print("Generator Configuration: ")
print(generator)
# create discriminator object:
dis_conf = get_config(args.discriminator_config)
dis_conf = list(map(get_layer, dis_conf.architecture))
discriminator = Discriminator(dis_conf)
print("Discriminator Configuration: ")
print(discriminator)
# create a gan from these
sagan = GAN(generator, discriminator, device=device)
# create optimizer for generator:
gen_optim = th.optim.Adam(filter(lambda p: p.requires_grad, generator.parameters()),
args.g_lr, [0, 0.9])
dis_optim = th.optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()),
args.d_lr, [0, 0.9])
# train the GAN
sagan.train(
data,
gen_optim,
dis_optim,
loss_fn=RelativisticAverageHingeGAN(device, discriminator),
num_epochs=args.num_epochs,
checkpoint_factor=args.checkpoint_factor,
data_percentage=args.data_percentage,
feedback_factor=31,
num_samples=64,
save_dir="models/relativistic/",
sample_dir="samples/4/",
log_dir="models/relativistic"
)
from attn_gan_pytorch.Networks import ConditionalGenerator
from attn_gan_pytorch.Utils import get_layer
from attn_gan_pytorch.ConfigManagement import get_config
from train import OneHotCIFAR10
# define the device for the training script
device = th.device("cuda" if th.cuda.is_available() else "cpu")
label = 9
num_samples = 12
latent_size = 128
num_classes = 10
config_file = "configs/dcgan-cifar-10-conditional/gen.conf"
model_file = "models/dcgan-cifar-10-conditional/GAN_GEN_12.pth"
gen_conf = get_config(config_file)
gen_conf = list(map(get_layer, gen_conf.architecture))
generator = ConditionalGenerator(gen_conf, 128).to(device)
generator.load_state_dict(th.load(model_file))
# generate one-hot samples:
labels = th.LongTensor([label for _ in range(num_samples)])
one_hot_labels = OneHotCIFAR10.one_hot_embedding(labels, num_classes)
# perform inference here:
expanded_labels = th.unsqueeze(th.unsqueeze(one_hot_labels, -1), -1)
latent_input = th.randn(num_samples, latent_size - num_classes, 1, 1)
gan_input = th.cat((latent_input, expanded_labels), dim=1).to(device)
generated_images = generator(gan_input).detach()
def main(args):
"""
Main function for the script
:param args: parsed command line arguments
:return: None
"""
from attn_gan_pytorch.Utils import get_layer
from attn_gan_pytorch.ConfigManagement import get_config
from attn_gan_pytorch.Networks import Generator, Discriminator, GAN
from data_processing.DataLoader import FlatDirectoryImageDataset, \
get_transform, get_data_loader
from attn_gan_pytorch.Losses import HingeGAN, RelativisticAverageHingeGAN
# create a data source:
celeba_dataset = FlatDirectoryImageDataset(args.images_dir,
transform=get_transform((64, 64)))
data = get_data_loader(celeba_dataset, args.batch_size, args.num_workers)
# create generator object:
gen_conf = get_config(args.generator_config)
gen_conf = list(map(get_layer, gen_conf.architecture))
generator = Generator(gen_conf, args.latent_size)
if args.generator_file is not None:
# load the weights into generator
generator.load_state_dict(th.load(args.generator_file))
print("Generator Configuration: ")
print(generator)
# create discriminator object:
dis_conf = get_config(args.discriminator_config)
dis_conf = list(map(get_layer, dis_conf.architecture))
discriminator = Discriminator(dis_conf)
if args.discriminator_file is not None:
# load the weights into discriminator
discriminator.load_state_dict(th.load(args.discriminator_file))
print("Discriminator Configuration: ")
print(discriminator)
# create a gan from these
fagan = GAN(generator, discriminator, device=device)
# create optimizer for generator:
gen_optim = th.optim.Adam(filter(lambda p: p.requires_grad, generator.parameters()),
args.g_lr, [0, 0.9])
dis_optim = th.optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()),
args.d_lr, [0, 0.9])
loss_name = args.loss_function.lower()
if loss_name == "hinge":
loss = HingeGAN
elif loss_name == "relativistic-hinge":
loss = RelativisticAverageHingeGAN
else:
raise Exception("Unknown loss function requested")
# train the GAN
fagan.train(
data,
gen_optim,
dis_optim,
loss_fn=loss(device, discriminator),
num_epochs=args.num_epochs,
checkpoint_factor=args.checkpoint_factor,
data_percentage=args.data_percentage,
feedback_factor=args.feedback_factor,
num_samples=64,
sample_dir=args.sample_dir,
save_dir=args.model_dir,
log_dir=args.model_dir,
start=args.start
)