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"
)
def main(args):
"""
Main function for the script
:param args: parsed command line arguments
:return: None
"""
from SMSG_GAN.SMSG_GAN import SMSG_GAN
from data_processing.DataLoader import FlatDirectoryImageDataset, \
get_transform, get_data_loader
from SMSG_GAN.Losses import HingeGAN, RelativisticAverageHingeGAN, \
StandardGAN, LSGAN
# create a data source:
celeba_dataset = FlatDirectoryImageDataset(
args.images_dir,
transform=get_transform(
(int(np.power(2, args.depth + 1)), int(np.power(2,
args.depth + 1)))))
data = get_data_loader(celeba_dataset, args.batch_size, args.num_workers)
# create a gan from these
smsg_gan = SMSG_GAN(depth=args.depth,
latent_size=args.latent_size,
device=device)
if args.generator_file is not None:
# load the weights into generator
smsg_gan.gen.load_state_dict(th.load(args.generator_file))
print("Generator Configuration: ")
print(smsg_gan.gen)
if args.discriminator_file is not None:
# load the weights into discriminator
smsg_gan.dis.load_state_dict(th.load(args.discriminator_file))
print("Discriminator Configuration: ")
print(smsg_gan.dis)
# create optimizer for generator:
gen_optim = th.optim.Adam(smsg_gan.gen.parameters(), args.g_lr, [0, 0.99])
dis_optim = th.optim.Adam(smsg_gan.dis.parameters(), args.d_lr, [0, 0.99])
loss_name = args.loss_function.lower()
if loss_name == "hinge":
loss = HingeGAN
elif loss_name == "relativistic-hinge":
loss = RelativisticAverageHingeGAN
elif loss_name == "standard-gan":
loss = StandardGAN
elif loss_name == "lsgan":
loss = LSGAN
else:
raise Exception("Unknown loss function requested")
# train the GAN
smsg_gan.train(data,
gen_optim,
dis_optim,
loss_fn=loss(device, smsg_gan.dis),
num_epochs=args.num_epochs,
checkpoint_factor=args.checkpoint_factor,
data_percentage=args.data_percentage,
feedback_factor=args.feedback_factor,
num_samples=args.num_samples,
sample_dir=args.sample_dir,
save_dir=args.model_dir,
log_dir=args.model_dir,
start=args.start)
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
)