def generate_samples(args, cc):
print("generating samples")
batch_size = 100
mixture_source = args.mixture_source
output_dir = args.output_dir
sample_size = args.sample_size
dataloader = cc.create_instance(cc.settings['dataloader']['dataset_name'])
network_factory = cc.create_instance(cc.settings['network']['name'],
dataloader.n_input_neurons)
population = Population(individuals=[], default_fitness=0)
mixture_definition = read_settings(
os.path.join(mixture_source, 'mixture.yml'))
for source, weight in mixture_definition.items():
path = os.path.join(mixture_source, source)
generator = network_factory.create_generator()
generator.net.load_state_dict(torch.load(path))
generator.net.eval()
population.individuals.append(
Individual(genome=generator, fitness=0, source=source))
dataset = MixedGeneratorDataset(
population, mixture_definition, sample_size * batch_size,
cc.settings['trainer']['mixture_generator_samples_mode'])
os.makedirs(output_dir, exist_ok=True)
LipizzanerMaster().save_samples(dataset, output_dir, dataloader,
sample_size, batch_size)
def calc_score(args, cc):
score_calc = ScoreCalculatorFactory.create()
cc.settings['general']['distribution']['client_id'] = 0
dataloader = cc.create_instance(cc.settings['dataloader']['dataset_name'])
network_factory = cc.create_instance(cc.settings['network']['name'],
dataloader.n_input_neurons)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
generator = network_factory.create_generator()
generator.net.load_state_dict(
torch.load(args.generator_file, map_location=device))
generator.net.eval()
individual = Individual(genome=generator, fitness=0, source='local')
dataset = MixedGeneratorDataset(
Population(individuals=[individual],
default_fitness=0), {'local': 1.0}, 50000,
cc.settings['trainer']['mixture_generator_samples_mode'])
output_dir = os.path.join(cc.output_dir, 'score')
os.makedirs(output_dir, exist_ok=True)
LipizzanerMaster().save_samples(dataset, output_dir, dataloader)
inc = score_calc.calculate(dataset)
_logger.info('Generator loaded from \'{}\' yielded a score of {}'.format(
args.generator_file, inc))
def calc_inception_score(args, cc):
inception_cal = InceptionCalculator(cuda=True)
dataloader = cc.create_instance(cc.settings['dataloader']['dataset_name'])
network_factory = cc.create_instance(cc.settings['network']['name'],
dataloader.n_input_neurons)
generator = network_factory.create_generator()
generator.net.load_state_dict(torch.load(args.inception_file))
generator.net.eval()
individual = Individual(genome=generator, fitness=0, source='local')
dataset = MixedGeneratorDataset(
Population(individuals=[individual], default_fitness=0),
{'local': 1.0}, 50000)
output_dir = os.path.join(cc.output_dir, 'inception_score')
os.makedirs(output_dir, exist_ok=True)
LipizzanerMaster().save_samples(dataset, output_dir, dataloader)
inc = inception_cal.calculate(dataset)
_logger.info(
'Generator loaded from \'{}\' yielded an inception score of {}'.format(
args.inception_file, inc))
if __name__ == '__main__':
os.environ['TORCH_MODEL_ZOO'] = os.path.join(os.getcwd(), 'output/.models')
parser = create_parser()
args = parser.parse_args(args=sys.argv[1:])
if 'cuda_device' in args and args.cuda_device:
print('Enforcing usage of CUDA device {}'.format(args.cuda_device))
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
if args.task == 'train':
if args.distributed:
if args.master:
initialize_settings(args)
LipizzanerMaster().run()
elif args.client:
LipizzanerClient().run()
else:
cc = initialize_settings(args)
lipizzaner = Lipizzaner()
lipizzaner.run(cc.settings['trainer']['n_iterations'])
elif args.task == 'score':
cc = initialize_settings(args)
calc_score(args, cc)
elif args.task == 'generate':
cc = initialize_settings(args)
generate_samples(args, cc)
else: