def run(args):
dataset = AlderleyWrapper()
save_args(args) # save command line to a file for reference
train_loader = alderley_cgan_data_loader(args, dataset=dataset) # get the data
model = CGANModel(
args,
discriminators=[PixelDiscriminator(), patchCDiscriminatorNetwork(args)],
generator=UnetUpsample())
# Build trainer
trainer = Trainer(model)
trainer.build_criterion(CWGANDiscriminatorLoss(penalty_weight=args.penalty_weight, model=model))
trainer.build_optimizer('Adam', [parameter for discriminator in model.discriminators for parameter in discriminator.parameters()], lr=args.discriminator_lr)
trainer.save_every((1, 'epochs'))
trainer.save_to_directory(args.save_directory)
trainer.set_max_num_epochs(args.epochs)
trainer.register_callback(CGenerateDataCallback(args, dataset=dataset))
trainer.register_callback(CGeneratorTrainingCallback(
args,
parameters=model.generator.parameters(),
criterion=WGANGeneratorLoss(), dataset=dataset))
trainer.bind_loader('train', train_loader, num_inputs=2)
# Custom logging configuration so it knows to log our images
if args.cuda:
trainer.cuda()
# Go!
trainer.fit()
def run(args):
dataset = RorschachWrapper()
save_args(args) # save command line to a file for reference
train_loader = rorschach_cgan_data_loader(args,
dataset=dataset) # get the data
# todo
model = patchCWGANModel(args,
discriminator=patchCDiscriminatorNetwork(args),
generator=CGeneratorNetwork(args))
# Build trainer
trainer = Trainer(model)
trainer.build_criterion(
CWGANDiscriminatorLoss(penalty_weight=args.penalty_weight,
model=model))
trainer.build_optimizer('Adam',
model.discriminator.parameters(),
lr=args.discriminator_lr)
trainer.save_every((1, 'epochs'))
trainer.save_to_directory(args.save_directory)
trainer.set_max_num_epochs(args.epochs)
trainer.register_callback(CGenerateDataCallback(args, dataset=dataset))
trainer.register_callback(
CGeneratorTrainingCallback(args,
parameters=model.generator.parameters(),
criterion=WGANGeneratorLoss(),
dataset=dataset))
trainer.bind_loader('train', train_loader, num_inputs=2)
# Custom logging configuration so it knows to log our images
logger = TensorboardLogger(log_scalars_every=(1, 'iteration'),
log_images_every=(args.log_image_frequency,
'iteration'))
trainer.build_logger(logger, log_directory=args.save_directory)
logger.observe_state('generated_images')
logger.observe_state('real_images')
logger._trainer_states_being_observed_while_training.remove(
'training_inputs')
if args.cuda:
trainer.cuda()
# Go!
trainer.fit()
