def __init__(self, cfg): super(LM, self).__init__() self.cfg = cfg # Dataset self.train_dataset = get_dataset( name=cfg.data.type, data_dir=cfg.data.train_dir, size=cfg.data.img_size, lsun_categories=cfg.data.lsun_categories_train ) # Number of labels self.nlabels = min(len(self.train_dataset.classes), cfg.data.nlabels) self.sample_nlabels = min(self.nlabels, cfg.train.sample_nlabels) # Create models self.generator = instantiate(cfg.generator) self.discriminator = instantiate(cfg.discriminator) self.gan_type = cfg.train.gan_type self.reg_type = cfg.train.reg_type self.reg_param = cfg.train.reg_param # Distributions self.ydist = get_ydist(self.nlabels) self.zdist = get_zdist(cfg.z_dist.type, cfg.z_dist.dim) # Save for tests ntest = cfg.train.batch_size x_real, self.ytest = utils.get_nsamples(self.train_dataloader(), ntest) self.ytest.clamp_(None, self.nlabels-1) self.ztest = self.zdist.sample((ntest,)) # Evaluator self.evaluator = Evaluator(self.generator, self.zdist, self.ydist, batch_size=cfg.train.batch_size)
# Create missing directories if not path.exists(out_dir): os.makedirs(out_dir) if not path.exists(checkpoint_dir): os.makedirs(checkpoint_dir) shutil.copyfile(sys.argv[0], out_dir + '/training_script.py') # Logger checkpoint_io = CheckpointIO(checkpoint_dir=checkpoint_dir) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Dataset train_dataset, nlabels = get_dataset( name=config['data']['type'], data_dir=config['data']['train_dir'], size=config['data']['img_size'], lsun_categories=config['data']['lsun_categories_train']) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, num_workers=config['training']['nworkers'], shuffle=True, pin_memory=True, sampler=None, drop_last=True) print('train_dataset=', train_dataset) test_dataset, _ = get_dataset( name=config['data']['type'], data_dir=config['data']['test_dir'], size=128,
def main(): pp = pprint.PrettyPrinter(indent=1) pp.pprint({ 'data': config['data'], 'generator': config['generator'], 'discriminator': config['discriminator'], 'clusterer': config['clusterer'], 'training': config['training'] }) is_cuda = torch.cuda.is_available() # Short hands batch_size = config['training']['batch_size'] log_every = config['training']['log_every'] inception_every = config['training']['inception_every'] backup_every = config['training']['backup_every'] sample_nlabels = config['training']['sample_nlabels'] nlabels = config['data']['nlabels'] sample_nlabels = min(nlabels, sample_nlabels) checkpoint_dir = path.join(out_dir, 'chkpts') nepochs = args.nepochs # Create missing directories if not path.exists(out_dir): os.makedirs(out_dir) if not path.exists(checkpoint_dir): os.makedirs(checkpoint_dir) # Logger checkpoint_io = CheckpointIO(checkpoint_dir=checkpoint_dir) device = torch.device("cuda:0" if is_cuda else "cpu") train_dataset, _ = get_dataset( name=config['data']['type'], data_dir=config['data']['train_dir'], size=config['data']['img_size'], deterministic=config['data']['deterministic']) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, num_workers=config['training']['nworkers'], shuffle=True, pin_memory=True, sampler=None, drop_last=True) # Create models generator, discriminator = build_models(config) # Put models on gpu if needed generator = generator.to(device) discriminator = discriminator.to(device) for name, module in discriminator.named_modules(): if isinstance(module, nn.Sigmoid): print('Found sigmoid layer in discriminator; not compatible with BCE with logits') exit() g_optimizer, d_optimizer = build_optimizers(generator, discriminator, config) devices = [int(x) for x in args.devices] generator = nn.DataParallel(generator, device_ids=devices) discriminator = nn.DataParallel(discriminator, device_ids=devices) # Register modules to checkpoint checkpoint_io.register_modules(generator=generator, discriminator=discriminator, g_optimizer=g_optimizer, d_optimizer=d_optimizer) # Logger logger = Logger(log_dir=path.join(out_dir, 'logs'), img_dir=path.join(out_dir, 'imgs'), monitoring=config['training']['monitoring'], monitoring_dir=path.join(out_dir, 'monitoring')) # Distributions ydist = get_ydist(nlabels, device=device) zdist = get_zdist(config['z_dist']['type'], config['z_dist']['dim'], device=device) ntest = config['training']['ntest'] x_test, y_test = utils.get_nsamples(train_loader, ntest) x_cluster, y_cluster = utils.get_nsamples(train_loader, config['clusterer']['nimgs']) x_test, y_test = x_test.to(device), y_test.to(device) z_test = zdist.sample((ntest, )) utils.save_images(x_test, path.join(out_dir, 'real.png')) logger.add_imgs(x_test, 'gt', 0) # Test generator if config['training']['take_model_average']: print('Taking model average') bad_modules = [nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d] for model in [generator, discriminator]: for name, module in model.named_modules(): for bad_module in bad_modules: if isinstance(module, bad_module): print('Batch norm in discriminator not compatible with exponential moving average') exit() generator_test = copy.deepcopy(generator) checkpoint_io.register_modules(generator_test=generator_test) else: generator_test = generator clusterer = get_clusterer(config)(discriminator=discriminator, x_cluster=x_cluster, x_labels=y_cluster, gt_nlabels=config['data']['nlabels'], **config['clusterer']['kwargs']) # Load checkpoint if it exists it = utils.get_most_recent(checkpoint_dir, 'model') if args.model_it == -1 else args.model_it it, epoch_idx, loaded_clusterer = checkpoint_io.load_models(it=it, load_samples='supervised' != config['clusterer']['name']) if loaded_clusterer is None: print('Initializing new clusterer. The first clustering can be quite slow.') clusterer.recluster(discriminator=discriminator) checkpoint_io.save_clusterer(clusterer, it=0) np.savez(os.path.join(checkpoint_dir, 'cluster_samples.npz'), x=x_cluster) else: print('Using loaded clusterer') clusterer = loaded_clusterer # Evaluator evaluator = Evaluator( generator_test, zdist, ydist, train_loader=train_loader, clusterer=clusterer, batch_size=batch_size, device=device, inception_nsamples=config['training']['inception_nsamples']) # Trainer trainer = Trainer(generator, discriminator, g_optimizer, d_optimizer, gan_type=config['training']['gan_type'], reg_type=config['training']['reg_type'], reg_param=config['training']['reg_param']) # Training loop print('Start training...') while it < args.nepochs * len(train_loader): epoch_idx += 1 for x_real, y in train_loader: it += 1 x_real, y = x_real.to(device), y.to(device) z = zdist.sample((batch_size, )) y = clusterer.get_labels(x_real, y).to(device) # Discriminator updates dloss, reg = trainer.discriminator_trainstep(x_real, y, z) logger.add('losses', 'discriminator', dloss, it=it) logger.add('losses', 'regularizer', reg, it=it) # Generators updates gloss = trainer.generator_trainstep(y, z) logger.add('losses', 'generator', gloss, it=it) if config['training']['take_model_average']: update_average(generator_test, generator, beta=config['training']['model_average_beta']) # Print stats if it % log_every == 0: g_loss_last = logger.get_last('losses', 'generator') d_loss_last = logger.get_last('losses', 'discriminator') d_reg_last = logger.get_last('losses', 'regularizer') print('[epoch %0d, it %4d] g_loss = %.4f, d_loss = %.4f, reg=%.4f' % (epoch_idx, it, g_loss_last, d_loss_last, d_reg_last)) if it % config['training']['recluster_every'] == 0 and it > config['training']['burnin_time']: # print cluster distribution for online methods if it % 100 == 0 and config['training']['recluster_every'] <= 100: print(f'[epoch {epoch_idx}, it {it}], distribution: {clusterer.get_label_distribution(x_real)}') clusterer.recluster(discriminator=discriminator, x_batch=x_real) # (i) Sample if necessary if it % config['training']['sample_every'] == 0: print('Creating samples...') x = evaluator.create_samples(z_test, y_test) x = evaluator.create_samples(z_test, clusterer.get_labels(x_test, y_test).to(device)) logger.add_imgs(x, 'all', it) for y_inst in range(sample_nlabels): x = evaluator.create_samples(z_test, y_inst) logger.add_imgs(x, '%04d' % y_inst, it) # (ii) Compute inception if necessary if it % inception_every == 0 and it > 0: print('PyTorch Inception score...') inception_mean, inception_std = evaluator.compute_inception_score() logger.add('metrics', 'pt_inception_mean', inception_mean, it=it) logger.add('metrics', 'pt_inception_stddev', inception_std, it=it) print(f'[epoch {epoch_idx}, it {it}] pt_inception_mean: {inception_mean}, pt_inception_stddev: {inception_std}') # (iii) Backup if necessary if it % backup_every == 0: print('Saving backup...') checkpoint_io.save('model_%08d.pt' % it, it=it) checkpoint_io.save_clusterer(clusterer, int(it)) logger.save_stats('stats_%08d.p' % it) if it > 0: checkpoint_io.save('model.pt', it=it)
os.makedirs(out_dir) if not path.exists(checkpoint_dir): os.makedirs(checkpoint_dir) # Logger checkpoint_io = CheckpointIO( checkpoint_dir=checkpoint_dir ) device = torch.device("cuda:0" if is_cuda else "cpu") # Dataset train_dataset, nlabels = get_dataset( name=config['data']['type'], data_dir=config['data']['train_dir'], size=config['data']['img_size'] ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, num_workers=config['training']['nworkers'], shuffle=True, pin_memory=True, sampler=None, drop_last=True ) # Number of labels nlabels = min(nlabels, config['data']['nlabels']) # Create models generator, discriminator = build_models(config)
def main(): checkpoint_dir = os.path.join(out_dir, 'chkpts') batch_size = config['training']['batch_size'] if 'cifar' in config['data']['train_dir'].lower(): name = 'cifar10' elif 'stacked_mnist' == config['data']['type']: name = 'stacked_mnist' else: name = 'image' if os.path.exists(os.path.join(out_dir, 'cluster_preds.npz')): # if we've already computed assignments, load them and move on with np.load(os.path.join(out_dir, 'cluster_preds.npz')) as f: y_reals = f['y_reals'] y_preds = f['y_preds'] else: train_dataset, _ = get_dataset(name=name, data_dir=config['data']['train_dir'], size=config['data']['img_size']) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, num_workers=config['training']['nworkers'], shuffle=True, pin_memory=True, sampler=None, drop_last=True) checkpoint_io = CheckpointIO(checkpoint_dir=checkpoint_dir) print('Loading clusterer:') most_recent = utils.get_most_recent( checkpoint_dir, 'model') if args.model_it is None else args.model_it clusterer = checkpoint_io.load_clusterer( most_recent, load_samples=False, pretrained=config['pretrained']) if isinstance(clusterer.discriminator, nn.DataParallel): clusterer.discriminator = clusterer.discriminator.module y_preds = [] y_reals = [] for batch_num, (x_real, y_real) in enumerate( tqdm(train_loader, total=len(train_loader))): y_pred = clusterer.get_labels(x_real.cuda(), None) y_preds.append(y_pred.detach().cpu()) y_reals.append(y_real) y_reals = torch.cat(y_reals).numpy() y_preds = torch.cat(y_preds).numpy() np.savez(os.path.join(out_dir, 'cluster_preds.npz'), y_reals=y_reals, y_preds=y_preds) if args.random: y_preds = np.random.randint(0, 100, size=y_reals.shape) nmi_score = nmi(y_preds, y_reals) purity = purity_score(y_preds, y_reals) print('nmi', nmi_score, 'purity', purity)
import numpy as np import os import sys sys.path.append(os.path.join(os.getcwd(), '..')) from gan_training.inputs import get_dataset torch.manual_seed(0) np.random.seed(0) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False if not os.path.exists('cifar_real'): os.makedirs('cifar_real') train_dataset, nlabels = get_dataset(name='cifar10', data_dir='../data', size=32) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=10000, num_workers=16, shuffle=True, pin_memory=True, sampler=None, drop_last=True) for x_real, y in train_loader: break x_real = x_real / 2 + 0.5
def main(): checkpoint_dir = os.path.join(out_dir, 'chkpts') most_recent = utils.get_most_recent( checkpoint_dir, 'model') if args.model_it is None else args.model_it cluster_path = os.path.join(out_dir, 'clusters') print('Saving clusters/samples to', cluster_path) os.makedirs(cluster_path, exist_ok=True) shutil.copyfile('seeing/lightbox.html', os.path.join(cluster_path, '+lightbox.html')) checkpoint_io = CheckpointIO(checkpoint_dir=checkpoint_dir) most_recent = utils.get_most_recent( checkpoint_dir, 'model') if args.model_it is None else args.model_it clusterer = checkpoint_io.load_clusterer(most_recent, pretrained=config['pretrained'], load_samples=False) if isinstance(clusterer.discriminator, nn.DataParallel): clusterer.discriminator = clusterer.discriminator.module model_path = os.path.join(checkpoint_dir, 'model_%08d.pt' % most_recent) sampler = SeededSampler(args.config, model_path=model_path, clusterer_path=os.path.join( checkpoint_dir, f'clusterer{most_recent}.pkl'), pretrained=config['pretrained']) if args.show_clusters: clusters = [[] for _ in range(config['generator']['nlabels'])] train_dataset, _ = get_dataset( name='webp' if 'cifar' not in config['data']['train_dir'].lower() else 'cifar10', data_dir=config['data']['train_dir'], size=config['data']['img_size']) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['training']['batch_size'], num_workers=config['training']['nworkers'], shuffle=True, pin_memory=True, sampler=None, drop_last=True) print('Generating clusters') for batch_num, (x_real, y_gt) in enumerate(train_loader): x_real = x_real.cuda() y_pred = clusterer.get_labels(x_real, y_gt) for i, yi in enumerate(y_pred): clusters[yi].append(x_real[i].cpu()) # don't generate too many, we're only visualizing 20 per cluster if batch_num * config['training']['batch_size'] >= 10000: break else: clusters = [None] * config['generator']['nlabels'] nimgs = 20 nrows = 4 for i in range(len(clusters)): if clusters[i] is None: pass elif len(clusters[i]) >= nimgs: cluster = torch.stack(clusters[i])[:nimgs] torchvision.utils.save_image(cluster * 0.5 + 0.5, os.path.join(cluster_path, f'{i}_real.png'), nrow=nrows) generated = [] for seed in range(nimgs): img = sampler.conditional_sample(i, seed=seed) generated.append(img.detach().cpu()) generated = torch.cat(generated) torchvision.utils.save_image(generated * 0.5 + 0.5, os.path.join(cluster_path, f'{i}_gen.png'), nrow=nrows) print('Clusters/samples can be visualized under', cluster_path)