if hparams["adversarial_loss"] == "nsgan":
adversarial_loss_fn = non_saturating_gan_losses
hparams["n_discriminator_iters"] = 1
elif hparams["adversarial_loss"] == "lsgan":
adversarial_loss_fn = lsgan_losses
hparams["n_discriminator_iters"] = 1
elif hparams["adversarial_loss"] == "wgan-gp":
adversarial_loss_fn = wgan_gp_losses
hparams["n_discriminator_iters"] = 5
hparams["wgan_gp_lambda"] = 10.0
else:
raise ValueError("Invalid adversarial loss fn.")
discogan = DiscoGAN(a_train=input_fn(args.x_train, hparams["batch_size"]),
a_test=input_fn(args.x_test, 3),
a_test_static=first_n(args.x_test, 10),
b_train=input_fn(args.y_train, hparams["batch_size"]),
b_test=input_fn(args.y_test, 3),
b_test_static=first_n(args.y_test, 10),
generator_fn=model_architecture.generator,
discriminator_fn=model_architecture.discriminator,
adversarial_loss_fn=adversarial_loss_fn,
**hparams)
experiments.run_experiment(model_dir=args.model_dir
or experiments.ROOT_RUNS_DIR /
experiments.experiment_name("discogan", hparams),
model=discogan,
n_training_step=args.steps)
if hparams["adversarial_loss"] == "lsgan":
adversarial_loss_fn = lsgan_losses
hparams["n_discriminator_iters"] = 1
elif hparams["adversarial_loss"] == "wgan-gp":
adversarial_loss_fn = wgan_gp_losses
hparams["n_discriminator_iters"] = 5
hparams["wgan_gp_lambda"] = 10.0
else:
raise ValueError("Invalid adversarial loss fn.")
cyclegan = CycleGAN(
a_train=input_fn(args.x_train, hparams["batch_size"]),
a_test=input_fn(args.x_test, 3),
a_test_static=first_n(args.x_test, 10),
b_train=input_fn(args.y_train, hparams["batch_size"]),
b_test=input_fn(args.y_test, 3),
b_test_static=first_n(args.y_test, 10),
generator_fn=model_architecture.generator,
discriminator_fn=model_architecture.discriminator,
adversarial_loss_fn=adversarial_loss_fn,
**hparams)
experiments.run_experiment(
model_dir=args.model_dir or experiments.ROOT_RUNS_DIR / experiments.experiment_name("cyclegan", hparams),
model=cyclegan,
n_training_step=args.steps)
elif hparams["adversarial_loss"] == "wgan-gp":
adversarial_loss_fn = wgan_gp_losses
hparams["n_discriminator_iters"] = 5
hparams["wgan_gp_lambda"] = 10.0
else:
raise ValueError("Invalid adversarial loss fn.")
attgan = AttGAN(
attribute_names=args.considered_attributes,
img=img_train,
attributes=attributes_train,
img_test=img_test,
attributes_test=attributes_test,
img_test_static=img_test_static,
attributes_test_static=attributes_test_static,
encoder_fn=model_architecture.encoder,
decoder_fn=model_architecture.decoder,
classifier_discriminator_shared_fn=model_architecture.
classifier_discriminator_shared,
classifier_private_fn=model_architecture.classifier_private,
discriminator_private_fn=model_architecture.discriminator_private,
adversarial_loss_fn=adversarial_loss_fn,
**hparams)
experiments.run_experiment(model_dir=args.model_dir
or experiments.ROOT_RUNS_DIR /
experiments.experiment_name("attgan", hparams),
model=attgan,
n_training_step=args.steps,
custom_init_op=init_op)
arg_parser.add_argument("--train_tfrecords",
nargs="+",
required=True,
help="train tfrecords files.")
arg_parser.add_argument("--test_tfrecords",
nargs="+",
required=True,
help="test tfrecords files.")
arg_parser.add_hparam("--batch_size",
default=32,
type=int,
help="Batch size")
args, hparams = arg_parser.parse_args()
ROOT_RUNS_DIR = Path("runs")
if args.model_dir is None:
model_dir = ROOT_RUNS_DIR / Path(
experiments.experiment_name("experimental", hparams))
else:
model_dir = Path(args.model_dir)
# TODO: Param for this. Handle mutual exclusiveness?
considered_attributes = [
"Smiling", "Male", "Mustache", "5_o_Clock_Shadow", "Blond_Hair"
]
run_training(model_dir, args.train_tfrecords, args.test_tfrecords,
considered_attributes, **hparams)
hparams["n_discriminator_iters"] = 1
elif hparams["adversarial_loss"] == "nsgan":
adversarial_loss_fn = non_saturating_gan_losses
hparams["n_discriminator_iters"] = 1
elif hparams["adversarial_loss"] == "wgan-gp":
adversarial_loss_fn = wgan_gp_losses
hparams["n_discriminator_iters"] = 5
hparams["wgan_gp_lambda"] = 10.0
else:
raise ValueError("Invalid adversarial loss fn.")
unit = UNIT(a_train=input_fn(args.x_train, hparams["batch_size"]),
a_test=input_fn(args.x_test, 3),
a_test_static=first_n(args.x_test, 10),
b_train=input_fn(args.y_train, hparams["batch_size"]),
b_test=input_fn(args.y_test, 3),
b_test_static=first_n(args.y_test, 10),
private_encoder_fn=model_architecture.encoder_private,
shared_encoder_fn=model_architecture.encoder_shared,
shared_decoder_fn=model_architecture.decoder_shared,
private_decoder_fn=model_architecture.decoder_private,
discriminator_fn=model_architecture.discriminator,
adversarial_loss_fn=adversarial_loss_fn,
**hparams)
experiments.run_experiment(model_dir=args.model_dir
or experiments.ROOT_RUNS_DIR /
experiments.experiment_name("unit", hparams),
model=unit,
n_training_step=args.steps)