def run(): parser = argparse.ArgumentParser(description="ALAE prepare SVHN") parser.add_argument( "--config-file", default="configs/svhn.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) random.seed(0) os.makedirs("SVHN", exist_ok=True) train = list(SVHN('.', split='train', download=True)) test = list(SVHN('.', split='test', download=True)) random.shuffle(train) svhn_images = np.stack([np.transpose(x[0], (2, 0, 1)) for x in train]) svhn_labels = np.stack([x[1] for x in train]) prepare_mnist(cfg, logger, svhn_images, svhn_labels, train=True) svhn_images = np.stack([np.transpose(x[0], (2, 0, 1)) for x in test]) svhn_labels = np.stack([x[1] for x in test]) prepare_mnist(cfg, logger, svhn_images, svhn_labels, train=False)
def train_net(args): torch.cuda.set_device(0) cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) torch.set_default_tensor_type('torch.cuda.FloatTensor') model = load_from('karras2019stylegan-ffhq-1024x1024.pkl', cfg) #model, Gs = load_from('karras2019stylegan-ffhq-1024x1024.pkl', cfg) # Generate image. #fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) #images = Gs.run(sample.cpu().detach().numpy(), None, truncation_psi=0.7, randomize_noise=True, output_transform=None) rnd = np.random.RandomState(5) latents = rnd.randn(1, cfg.MODEL.LATENT_SPACE_SIZE) sample = torch.tensor(latents).float().cuda() save_sample( model, sample, ) #png_filename = os.path.join('example.png') #PIL.Image.fromarray(images[0], 'RGB').save(png_filename) model_dict = { 'generator_s': model.generator, 'mapping_s': model.mapping, 'dlatent_avg': model.dlatent_avg, } checkpointer = Checkpointer(cfg, model_dict, logger=logger, save=True) checkpointer.save('karras2019stylegan-ffhq')
def run(fn): parser = argparse.ArgumentParser( description="PETScanner photon detector training scripts") parser.add_argument( "--config-file", default="configs/experiment.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() torch.cuda.set_device(0) cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) fh = logging.FileHandler(os.path.join(output_dir, 'log.txt')) fh.setLevel(logging.DEBUG) fh.setFormatter(formatter) logger.addHandler(fh) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) torch.set_default_tensor_type('torch.cuda.FloatTensor') device = torch.cuda.current_device() print("Running on ", torch.cuda.get_device_name(device)) fn(cfg, logger)
def create_model(self, config_path="/configs/ffhq.yaml"): self.cfg = get_cfg_defaults() self.cfg.merge_from_file(REPO_PATH + config_path) logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) logging.basicConfig( stream=sys.stdout, level=logging.INFO) # print out all the info messages # print("cfg", cfg) model = Model(startf=self.cfg.MODEL.START_CHANNEL_COUNT, layer_count=self.cfg.MODEL.LAYER_COUNT, maxf=self.cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=self.cfg.MODEL.LATENT_SPACE_SIZE, truncation_psi=self.cfg.MODEL.TRUNCATIOM_PSI, truncation_cutoff=self.cfg.MODEL.TRUNCATIOM_CUTOFF, mapping_layers=self.cfg.MODEL.MAPPING_LAYERS, channels=self.cfg.MODEL.CHANNELS, generator=self.cfg.MODEL.GENERATOR, encoder=self.cfg.MODEL.ENCODER) model.cuda() model.eval() model.requires_grad_(False) decoder = model.decoder encoder = model.encoder mapping_tl = model.mapping_tl mapping_fl = model.mapping_fl dlatent_avg = model.dlatent_avg logger.info("Trainable parameters generator:") print(count_parameters(decoder)) logger.info("Trainable parameters discriminator:") print(count_parameters(encoder)) model_dict = { 'discriminator_s': encoder, 'generator_s': decoder, 'mapping_tl_s': mapping_tl, 'mapping_fl_s': mapping_fl, 'dlatent_avg': dlatent_avg } checkpointer = Checkpointer(self.cfg, model_dict, {}, logger=logger, save=False) self.checkpointer = checkpointer self.model = model self.layer_count = self.cfg.MODEL.LAYER_COUNT self.encoder = encoder self.decoder = decoder
def run(): parser = argparse.ArgumentParser(description="ALAE. prepare cifar10") parser.add_argument( "--config-file", default="configs/cifar10.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) random.seed(0) dlutils.download.cifar10() cifar10 = dlutils.reader.Cifar10('cifar10/cifar-10-batches-bin', train=True, test=False).items random.shuffle(cifar10) cifar10_images = np.stack([x[1] for x in cifar10]) cifar10_labels = np.stack([x[0] for x in cifar10]) prepare_cifar10(cfg, logger, cifar10_images, cifar10_labels, train=False) prepare_cifar10(cfg, logger, cifar10_images, cifar10_labels, train=True)
def run(): parser = argparse.ArgumentParser( description="ALAE. Split FFHQ into parts for training and testing") parser.add_argument( "--config-file", default="configs/ffhq.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) fh = logging.FileHandler(os.path.join(output_dir, 'log.txt')) fh.setLevel(logging.DEBUG) fh.setFormatter(formatter) logger.addHandler(fh) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) split_tfrecord(cfg, logger)
def __init__(self, folder, device=None, remove=False, resume=False, run_prefix='run', tensorboard=True): super(ExperimentConfig, self).__init__() self.folder = folder if remove: os.system(f'rm -r {folder}/{run_prefix}*') self.files = os.popen(f'ls {folder}').read().split('\n') max_run = 0 for f in self.files: match = re.search(f'^{run_prefix}(\d+)$', f) if match: max_run = int(match[1]) if not resume: max_run += 1 log_dir = f'{folder}/{run_prefix}{max_run}' if tensorboard: self.writer = SummaryWriter(log_dir=log_dir, comment=log_dir) self.cfg = get_cfg_defaults() self.cfg.merge_from_file(f'{folder}/config.yml') self.cfg.freeze() for k in valid_values: if self.cfg[k] not in valid_values[k]: raise Exception(f"Invalid value for {k}") for k in self.cfg: setattr(self, k, self.cfg[k]) self.gpu_id = None if device is not None: self.device = torch.device(device) matches = re.match('.*:(\d+)', device) if matches: self.gpu_id = int(matches[1]) else: self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu")
def load_file(file_loc): cfg = _C.clone() cfg.merge_from_file(file_loc) if cfg.INHERIT != '': base_cfg = load_file(cfg.INHERIT) base_cfg.merge_from_file(file_loc) cfg = base_cfg # special sub config if len(cfg.MODEL_CONFIG_LOCATION) > 0: sub_cfg = get_cfg_defaults() sub_cfg.merge_from_file(cfg.MODEL_CONFIG_LOCATION + "/config.yml") cfg.defrost() cfg.MODEL_CONFIG = sub_cfg cfg.freeze() return cfg
def run(): parser = argparse.ArgumentParser( description="Adversarial, hierarchical style VAE") parser.add_argument( "--config-file", default="configs/afhq.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) prepare_afhq(cfg, logger, True) prepare_afhq(cfg, logger, False)
def convert(args): #torch.cuda.set_device(0) cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) logger.info(args) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) #torch.set_default_tensor_type('torch.cuda.FloatTensor') #model, Gs = load_from('karras2019stylegan-ffhq-1024x1024.pkl', cfg) model, _ = load_from('./pre-model/karras2019stylegan-cars-512x384.pkl', cfg) model_dict = { 'generator_s': model.generator, 'mapping_fl_s': model.mapping, 'dlatent_avg': model.dlatent_avg, } checkpointer = Checkpointer(cfg, model_dict, logger=logger, save=True) #checkpointer.save('karras2019stylegan-ffhq') checkpointer.save('karras2019stylegan-cars')
def train_net(gpu_id, args): torch.cuda.set_device(0) cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.freeze() logger = logging.getLogger("logger") logger.setLevel(logging.DEBUG) output_dir = cfg.OUTPUT_DIR os.makedirs(output_dir, exist_ok=True) if gpu_id == 0: ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( "%(asctime)s %(name)s %(levelname)s: %(message)s") ch.setFormatter(formatter) logger.addHandler(ch) fh = logging.FileHandler(os.path.join(output_dir, 'log.txt')) fh.setLevel(logging.DEBUG) fh.setFormatter(formatter) logger.addHandler(fh) logger.info(args) logger.info("Using {} GPUs".format(1)) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) torch.set_default_tensor_type('torch.cuda.FloatTensor') train(cfg, logger)
'dlatent_avg': dlatent_avg } checkpointer = Checkpointer(cfg, model_dict, {}, logger=logger, save=False) checkpointer.load() model.eval() layer_count = cfg.MODEL.LAYER_COUNT logger.info("Generating...") decoder = nn.DataParallel(decoder) mapping_fl = nn.DataParallel(mapping_fl) with torch.no_grad(): gen = ImageGenerator(cfg, num_samples=60000, minibatch_gpu=8) gen.evaluate(logger, mapping_fl, decoder, cfg.DATASET.MAX_RESOLUTION_LEVEL - 2) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-generate-images-for-attribute-classifications', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
save_image(src_originals[i] * 0.5 + 0.5, './style_mixing/output/%s/source_%d.png' % (cfg.NAME, i)) place(canvas, src_originals[i], 1 + i, 0) for i in range(dst_len): save_image(dst_originals[i] * 0.5 + 0.5, './style_mixing/output/%s/dst_coarse_%d.png' % (cfg.NAME, i)) place(canvas, dst_originals[i], 0, 1 + i) style_ranges = [range(0, 4)] * 3 + [range(4, 8)] * 2 + [range(8, layer_count * 2)] def mix_styles(style_src, style_dst, r): style = style_dst.clone() style[:, r] = style_src[:, r] return style for row in range(dst_len): row_latents = torch.stack([dst_latents[row]] * src_len) style = mix_styles(src_latents, row_latents, style_ranges[row]) rec = model.decoder(style, layer_count - 1, 1, noise=True) for j in range(rec.shape[0]): save_image(rec[j] * 0.5 + 0.5, './style_mixing/output/%s/rec_coarse_%d_%d.png' % (cfg.NAME, row, j)) place(canvas, rec[j], 1 + j, 1 + row) save_image(torch.Tensor(canvas), './style_mixing/output/%s/stylemix.png' % cfg.NAME) if __name__ == "__main__": gpu_count = 1 run(main, get_cfg_defaults(), description='SandwichStyleVAE-style-mixing', default_config='./configs/celeba-hq256-generate.yaml', world_size=gpu_count, write_log=False)
def main(folding_id, inliner_classes, ic, total_classes, mul, folds=5): cfg = get_cfg_defaults() cfg.merge_from_file('configs/mnist.yaml') cfg.freeze() logger = logging.getLogger("logger") torch.set_default_tensor_type('torch.cuda.FloatTensor') device = torch.cuda.current_device() print("Running on ", torch.cuda.get_device_name(device)) train_set, valid_set, test_set = make_datasets(cfg, folding_id, inliner_classes) print('Validation set size: %d' % len(valid_set)) print('Test set size: %d' % len(test_set)) train_set.shuffle() G = Generator(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) E = Encoder(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) G.load_state_dict(torch.load("models/Gmodel_%d_%d.pkl" %(folding_id, ic))) E.load_state_dict(torch.load("models/Emodel_%d_%d.pkl" %(folding_id, ic))) G.eval() E.eval() sample = torch.randn(64, cfg.MODEL.LATENT_SIZE).to(device) sample = G(sample.view(-1, cfg.MODEL.LATENT_SIZE, 1, 1)).cpu() save_image(sample.view(64, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE), 'sample.png') counts, bin_edges, gennorm_param = extract_statistics(cfg, train_set, inliner_classes, E, G) def run_novely_prediction_on_dataset(dataset, percentage, concervative=False): dataset.shuffle() dataset = create_set_with_outlier_percentage(dataset, inliner_classes, percentage, concervative) result = [] gt_novel = [] data_loader = make_dataloader(dataset, cfg.TEST.BATCH_SIZE, torch.cuda.current_device()) include_jacobian = True N = (cfg.MODEL.INPUT_IMAGE_SIZE * cfg.MODEL.INPUT_IMAGE_SIZE - cfg.MODEL.LATENT_SIZE) * mul logC = loggamma(N / 2.0) - (N / 2.0) * np.log(2.0 * np.pi) def logPe_func(x): # p_{\|W^{\perp}\|} (\|w^{\perp}\|) # \| w^{\perp} \|}^{m-n} return logC - (N - 1) * np.log(x) + np.log(r_pdf(x, bin_edges, counts)) for label, x in data_loader: x = x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS * cfg.MODEL.INPUT_IMAGE_SIZE * cfg.MODEL.INPUT_IMAGE_SIZE) x = Variable(x.data, requires_grad=True) z,logvar,mu = E(x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE)) recon_batch = G(z) z = z.squeeze() if include_jacobian: J = compute_jacobian(x, z) J = J.cpu().numpy() z = z.cpu().detach().numpy() recon_batch = recon_batch.squeeze().cpu().detach().numpy() x = x.squeeze().cpu().detach().numpy() for i in range(x.shape[0]): if include_jacobian: u, s, vh = np.linalg.svd(J[i, :, :], full_matrices=False) logD = -np.sum(np.log(np.abs(s))) # | \mathrm{det} S^{-1} | # logD = np.log(np.abs(1.0/(np.prod(s)))) else: logD = 0 p = scipy.stats.gennorm.pdf(z[i], gennorm_param[0, :], gennorm_param[1, :], gennorm_param[2, :]) logPz = np.sum(np.log(p)) # Sometimes, due to rounding some element in p may be zero resulting in Inf in logPz # In this case, just assign some large negative value to make sure that the sample # is classified as unknown. if not np.isfinite(logPz): logPz = -1000 distance = np.linalg.norm(x[i].flatten() - recon_batch[i].flatten()) logPe = logPe_func(distance) P = logD + logPz + logPe result.append(P) gt_novel.append(label[i].item() in inliner_classes) result = np.asarray(result, dtype=np.float32) ground_truth = np.asarray(gt_novel, dtype=np.float32) return result, ground_truth def compute_threshold(valid_set, percentage): y_scores, y_true = run_novely_prediction_on_dataset(valid_set, percentage, concervative=True) minP = min(y_scores) - 1 maxP = max(y_scores) + 1 y_false = np.logical_not(y_true) def evaluate(e): y = np.greater(y_scores, e) true_positive = np.sum(np.logical_and(y, y_true)) false_positive = np.sum(np.logical_and(y, y_false)) false_negative = np.sum(np.logical_and(np.logical_not(y), y_true)) return get_f1(true_positive, false_positive, false_negative) best_th, best_f1 = find_maximum(evaluate, minP, maxP, 1e-4) logger.info("Best e: %f best f1: %f" % (best_th, best_f1)) return best_th def test(test_set, percentage, threshold): y_scores, y_true = run_novely_prediction_on_dataset(test_set, percentage, concervative=True) return evaluate(logger, percentage, inliner_classes, y_scores, threshold, y_true) percentages = cfg.DATASET.PERCENTAGES # percentages = [50] results = {} for p in percentages: plt.figure(num=None, figsize=(8, 6), dpi=180, facecolor='w', edgecolor='k') e = compute_threshold(valid_set, p) results[p] = test(test_set, p, e) return results
with torch.no_grad(): for img in sample: x = torch.tensor(np.asarray(img, dtype=np.float32), device='cpu', requires_grad=True).cuda() / 127.5 - 1. if x.shape[0] == 4: x = x[:3] latents = encode(x[None, ...].cuda()) f = decode(latents) r = torch.cat([x[None, ...].detach().cpu(), f.detach().cpu()], dim=3) canvas.append(r) return canvas sample = next(b) canvas = make(sample) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, './make_figures/reconstructions_ffhq_real_1.png', nrow=2, pad_value=1.0) sample = next(b) canvas = make(sample) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, './make_figures/reconstructions_ffhq_real_2.png', nrow=2, pad_value=1.0) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='SoftIntroVAE-reconstruction-ffhq', default_config='./configs/ffhq256.yaml', world_size=gpu_count, write_log=False)
for i in range(src_len): save_image(src_originals[i] * 0.5 + 0.5, 'style_mixing/output/%s/source_%d.png' % (cfg.NAME, i)) place(canvas, src_originals[i], 1 + i, 0) for i in range(dst_len): save_image(dst_originals[i] * 0.5 + 0.5, 'style_mixing/output/%s/dst_coarse_%d.png' % (cfg.NAME, i)) place(canvas, dst_originals[i], 0, 1 + i) style_ranges = [range(0, 4)] * 3 + [range(4, 8)] * 2 + [range(8, layer_count * 2)] def mix_styles(style_src, style_dst, r): style = style_dst.clone() style[:, r] = style_src[:, r] return style for row in range(dst_len): row_latents = torch.stack([dst_latents[row]] * src_len) style = mix_styles(src_latents, row_latents, style_ranges[row]) rec = model.decoder(style, layer_count - 1, 1, noise=True) for j in range(rec.shape[0]): save_image(rec[j] * 0.5 + 0.5, 'style_mixing/output/%s/rec_coarse_%d_%d.png' % (cfg.NAME, row, j)) place(canvas, rec[j], 1 + j, 1 + row) save_image(torch.Tensor(canvas), 'style_mixing/output/%s/stylemix.png' % cfg.NAME) if __name__ == "__main__": gpu_count = 1 run(main, get_cfg_defaults(), description='ALAE-style-mixing', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
factor = x.shape[2] // im_size if factor != 1: x = torch.nn.functional.avg_pool2d(x[None, ...], factor, factor)[0] assert x.shape[2] == im_size latents = encode(x[None, ...].cuda()) f = decode(latents) r = torch.cat([x[None, ...].detach().cpu(), f.detach().cpu()], dim=3) canvas.append(r) return canvas def chunker_list(seq, n): return [seq[i * n:(i + 1) * n] for i in range((len(seq) + n - 1) // n)] paths = chunker_list(paths, 8 * 3) for i, chunk in enumerate(paths): canvas = make(chunk) canvas = torch.cat(canvas, dim=0) save_path = 'make_figures/output/%s/reconstructions_%d.png' % (cfg.NAME, i) os.makedirs(os.path.dirname(save_path), exist_ok=True) save_image(canvas * 0.5 + 0.5, save_path, nrow=3, pad_value=1.0) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-figure-reconstructions-paged', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
def main(): parser = argparse.ArgumentParser(description="Process training arguments.") parser.add_argument('--config', type=str, default="configurations/ppo_baseline_cuda.yaml", help="config file name (located in config dir)") args = parser.parse_args() # create configuration cfg = get_cfg_defaults() cfg.merge_from_file(args.config) print(cfg.TRAIN.TOTAL_TIMESTEPS) # create experiment directory exp_dir = f"runs/{cfg.EXPERIMENT_NAME}/{datetime.now().strftime('%Y-%m-%d-%H-%M')}" os.makedirs(exp_dir, exist_ok=True) # create logger format_strs = ['csv', 'stdout'] logger.configure(dir=exp_dir, format_strs=format_strs, log_suffix=datetime.now().strftime('%Y-%m-%d-%H-%M')) # create (vectorized) procgen environment logger.info("creating environment") venv = ProcgenEnv(num_envs=cfg.TRAIN.NUM_ENVS, env_name="fruitbot", num_levels=cfg.TRAIN.NUM_LEVELS, start_level=cfg.TRAIN.LEVEL_SEED, distribution_mode="easy") venv = VecExtractDictObs(venv, "rgb") venv = VecMonitor( venv=venv, filename=None, keep_buf=100, ) venv = VecNormalize(venv=venv, ob=False) test_venv = ProcgenEnv(num_envs=cfg.TEST.NUM_ENVS, env_name="fruitbot", num_levels=cfg.TEST.NUM_LEVELS, start_level=cfg.TEST.LEVEL_SEED, distribution_mode="easy") test_venv = VecExtractDictObs(test_venv, "rgb") test_venv = VecMonitor( venv=test_venv, filename=None, keep_buf=100, ) test_venv = VecNormalize(venv=test_venv, ob=False) # create tensorflow session logger.info("creating tf session") config = tf.ConfigProto() config.gpu_options.allow_growth = True # pylint: disable=E1101 sess = tf.Session(config=config) sess.__enter__() # create cnn todo: make this less ugly conv_fn = None logger.info("building cnn") if cfg.TRAIN.NETWORK == "NATURE_CNN": conv_fn = lambda x: nature_cnn(x) elif cfg.TRAIN.NETWORK == "IMPALA_CNN": conv_fn = lambda x: build_impala_cnn( x, depths=[16, 32, 32], emb_size=256) # training logger.info("training") if cfg.TRAIN.POLICY == "A2C": a2c.learn(env=venv, network=conv_fn, total_timesteps=cfg.TRAIN.TOTAL_TIMESTEPS, nsteps=cfg.TRAIN.BATCH_SIZE, log_interval=1, eval_env=test_venv, augment=cfg.TRAIN.AUGMENT) elif cfg.TRAIN.POLICY == "ACKTR": acktr.learn(env=venv, network=conv_fn, total_timesteps=cfg.TRAIN.TOTAL_TIMESTEPS, nsteps=cfg.TRAIN.BATCH_SIZE, log_interval=1, eval_env=test_venv, augment=cfg.TRAIN.AUGMENT, seed=None) elif cfg.TRAIN.POLICY == "PPO": ppo2.learn(env=venv, eval_env=test_venv, network=conv_fn, total_timesteps=cfg.TRAIN.TOTAL_TIMESTEPS, save_interval=5, nsteps=cfg.TRAIN.BATCH_SIZE, nminibatches=cfg.TRAIN.MINIBATCHES, lam=cfg.TRAIN.LAM, gamma=cfg.TRAIN.GAMMA, noptepochs=cfg.TRAIN.NUM_EPOCHS, log_interval=1, clip_vf=cfg.TRAIN.USE_VF_CLIPPING, lr=cfg.TRAIN.LR, cliprange=cfg.TRAIN.CLIP_RANGE, update_fn=None, init_fn=None, vf_coef=0.5, max_grad_norm=0.5, augment=cfg.TRAIN.AUGMENT, load_path=cfg.TRAIN.PRETRAINED)
lod2batch.step() # if lod2batch.is_time_to_save(): # checkpointer.save("model_tmp_intermediate_lod%d" % lod_for_saving_model) if lod2batch.is_time_to_report(): save_sample(lod2batch, tracker, sample, samplez, x, logger, model_s, cfg, encoder_optimizer, generator_optimizer, output_folder) scheduler.step() if epoch % 20 == 0: save(epoch) save_sample(lod2batch, tracker, sample, samplez, x, logger, model_s, cfg, encoder_optimizer, generator_optimizer, output_folder) logger.info("Training finish!... save training results") if epoch is not None: save(epoch) best_model_name, best_model_score = scores_list[0] for model_name, model_score in scores_list: if model_score >= best_model_score: best_model_name, best_model_score = model_name, model_score checkpointer.tag_best_checkpoint(best_model_name) if __name__ == "__main__": gpu_count = torch.cuda.device_count() run(train, get_cfg_defaults(), description='', default_config='configs/mnist_os.yaml', world_size=gpu_count)
im = img.transpose((2, 0, 1)) x = torch.tensor(np.asarray(im, dtype=np.float32), device='cpu', requires_grad=True).cuda() / 127.5 - 1. if x.shape[0] == 4: x = x[:3] while x.shape[2] != model.decoder.layer_to_resolution[6]: x = F.avg_pool2d(x, 2, 2) latents = encode(x[None, ...].cuda()) f = decode(latents) r = torch.cat([x[None, ...].detach().cpu(), f.detach().cpu()], dim=3) canvas.append(r) return canvas canvas = make(paths[:10]) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, 'make_figures/output/reconstructions_celeba_1.png', nrow=2, pad_value=1.0) canvas = make(paths[10:20]) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, 'make_figures/output/reconstructions_celeba_2.png', nrow=2, pad_value=1.0) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-reconstruction-bedroom', default_config='configs/celeba-hq256.yaml', world_size=gpu_count, write_log=False)
def make(sample): canvas = [] with torch.no_grad(): for img in sample: x = torch.tensor(np.asarray(img, dtype=np.float32), device='cpu', requires_grad=True).cuda() / 127.5 - 1. if x.shape[0] == 4: x = x[:3] latents = encode(x[None, ...].cuda()) f = decode(latents) r = torch.cat([x[None, ...].detach().cpu(), f.detach().cpu()], dim=3) canvas.append(r) return canvas sample = next(b) canvas = make(sample) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, 'make_figures/reconstructions_ffhq_real_1.png', nrow=2, pad_value=1.0) sample = next(b) canvas = make(sample) canvas = torch.cat(canvas, dim=0) save_image(canvas * 0.5 + 0.5, 'make_figures/reconstructions_ffhq_real_2.png', nrow=2, pad_value=1.0) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-reconstruction-ffhq', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
def make(paths): with torch.no_grad(): for filename in paths: img = np.asarray(Image.open(path + '/' + filename)) if img.shape[2] == 4: img = img[:, :, :3] im = img.transpose((2, 0, 1)) x = torch.tensor(np.asarray(im, dtype=np.float32), device='cpu', requires_grad=True).cuda() / 127.5 - 1. if x.shape[0] == 4: x = x[:3] while x.shape[2] != model.decoder.layer_to_resolution[6]: x = F.avg_pool2d(x, 2, 2) latents = encode(x[None, ...].cuda()) f = decode(latents) r = torch.cat([x[None, ...].detach().cpu(), f.detach().cpu()], dim=3) os.makedirs('make_figures/output/pioneer/', exist_ok=True) save_image(f.detach().cpu() * 0.5 + 0.5, 'make_figures/output/pioneer/%s_alae.png' % filename[:-9], nrow=1, pad_value=1.0) make(paths) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-reconstructions-celeb-hq256-on-pioneer-examples', default_config='configs/celeba-hq256.yaml', world_size=gpu_count, write_log=False)
for i in range(height): for j in range(width): kv = i / (height - 1.0) kh = j / (width - 1.0) ka = (1.0 - kh) * (1.0 - kv) kb = kh * (1.0 - kv) kc = (1.0 - kh) * kv kd = kh * kv w = ka * wa + kb * wb + kc * wc + kd * wd interpolated = make(w) images.append(interpolated) images = torch.cat(images) path = './make_figures/output' os.makedirs(path, exist_ok=True) os.makedirs(os.path.join(path, cfg.NAME), exist_ok=True) save_image(images * 0.5 + 0.5, './make_figures/output/%s/interpolations.png' % cfg.NAME, nrow=width) save_image(images * 0.5 + 0.5, './make_figures/output/%s/interpolations.jpg' % cfg.NAME, nrow=width) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='SoftIntroVAE-interpolations', default_config='./configs/ffhq256.yaml', world_size=gpu_count, write_log=False)
height = 7 width = 7 images = [] for i in range(height): for j in range(width): kv = i / (height - 1.0) kh = j / (width - 1.0) ka = (1.0 - kh) * (1.0 - kv) kb = kh * (1.0 - kv) kc = (1.0 - kh) * kv kd = kh * kv w = ka * wa + kb * wb + kc * wc + kd * wd interpolated = make(w) images.append(interpolated) images = torch.cat(images) save_image(images * 0.5 + 0.5, 'make_figures/output/%s/interpolations.png' % cfg.NAME, nrow=width) save_image(images * 0.5 + 0.5, 'make_figures/output/%s/interpolations.jpg' % cfg.NAME, nrow=width) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-interpolations', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
'dlatent_avg': dlatent_avg } checkpointer = Checkpointer(cfg, model_dict, {}, logger=logger, save=False) checkpointer.load() model.eval() layer_count = cfg.MODEL.LAYER_COUNT logger.info("Extracting attributes") decoder = nn.DataParallel(decoder) indices = [0, 1, 2, 3, 4, 10, 11, 17, 19] with torch.no_grad(): p = Predictions(cfg, minibatch_gpu=4) for i in indices: p.evaluate(logger, mapping_fl, decoder, cfg.DATASET.MAX_RESOLUTION_LEVEL - 2, i) if __name__ == "__main__": gpu_count = 1 run(main, get_cfg_defaults(), description='StyleGAN', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
def main(): # get model path parser = argparse.ArgumentParser(description="Parse testing arguments") parser.add_argument('--model_path', type=str, default=None, help='Path to model checkpoint.') parser.add_argument('--config', type=str, default='configurations/ppo_baseline_cuda.yaml', help='Path to configuration file.') args = parser.parse_args() if args.model_path is None or not os.path.exists(args.model_path): raise OSError("Invalid model file supplied") # create configuration cfg = get_cfg_defaults() cfg.merge_from_file(args.config) # create save directory model_file_path = args.model_path exp_creation_time = os.path.normpath(model_file_path).split(os.sep)[-3] print(exp_creation_time) exp_dir = f"runs/{cfg.EXPERIMENT_NAME}/{exp_creation_time}_test/" os.makedirs(exp_dir, exist_ok=True) # create logger format_strs = ['csv', 'stdout'] logger.configure(dir=exp_dir, format_strs=format_strs, log_suffix=datetime.now().strftime('%Y-%m-%d-%H-%M')) # create (vectorized) procgen environment logger.info("creating environment") venv = ProcgenEnv(num_envs=cfg.TEST.NUM_ENVS, env_name="fruitbot", num_levels=cfg.TEST.NUM_LEVELS, start_level=cfg.TEST.LEVEL_SEED, distribution_mode="easy") venv = VecExtractDictObs(venv, "rgb") venv = VecMonitor( venv=venv, filename=None, keep_buf=100, ) venv = VecNormalize(venv=venv, ob=False) # create tensorflow session logger.info("creating tf session") config = tf.ConfigProto() config.gpu_options.allow_growth = True # pylint: disable=E1101 sess = tf.Session(config=config) sess.__enter__() # create cnn todo: make this less ugly conv_fn = None logger.info("building cnn") if cfg.TRAIN.NETWORK == "NATURE_CNN": conv_fn = lambda x: nature_cnn(x) elif cfg.TRAIN.NETWORK == "IMPALA_CNN": conv_fn = lambda x: build_impala_cnn( x, depths=[16, 32, 32], emb_size=256) # training logger.info("testing") ppo2.learn(env=venv, network=conv_fn, total_timesteps=cfg.TEST.TIMESTEPS, save_interval=0, nsteps=cfg.TEST.BATCH_SIZE, nminibatches=cfg.TRAIN.MINIBATCHES, lam=cfg.TRAIN.LAM, gamma=cfg.TRAIN.GAMMA, noptepochs=cfg.TRAIN.NUM_EPOCHS, log_interval=1, clip_vf=cfg.TRAIN.USE_VF_CLIPPING, lr=cfg.TRAIN.LR, cliprange=cfg.TRAIN.CLIP_RANGE, update_fn=None, init_fn=None, vf_coef=0.5, max_grad_norm=0.5, test=True, load_path=model_file_path)
def train(folding_id, inliner_classes, ic): cfg = get_cfg_defaults() cfg.merge_from_file('configs/mnist.yaml') cfg.freeze() logger = logging.getLogger("logger") zsize = cfg.MODEL.LATENT_SIZE output_folder = os.path.join('results_' + str(folding_id) + "_" + "_".join([str(x) for x in inliner_classes])) os.makedirs(output_folder, exist_ok=True) os.makedirs('models', exist_ok=True) train_set, _, _ = make_datasets(cfg, folding_id, inliner_classes) logger.info("Train set size: %d" % len(train_set)) G = Generator(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) G.weight_init(mean=0, std=0.02) D = Discriminator(channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) D.weight_init(mean=0, std=0.02) E = Encoder(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS) E.weight_init(mean=0, std=0.02) if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH: ZD = ZDiscriminator_mergebatch(zsize, cfg.TRAIN.BATCH_SIZE) else: ZD = ZDiscriminator(zsize, cfg.TRAIN.BATCH_SIZE) ZD.weight_init(mean=0, std=0.02) lr = cfg.TRAIN.BASE_LEARNING_RATE G_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999)) D_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999)) GE_optimizer = optim.Adam(list(E.parameters()) + list(G.parameters()), lr=lr, betas=(0.5, 0.999)) ZD_optimizer = optim.Adam(ZD.parameters(), lr=lr, betas=(0.5, 0.999)) BCE_loss = nn.BCELoss() sample = torch.randn(64, zsize).view(-1, zsize, 1, 1) tracker = LossTracker(output_folder=output_folder) for epoch in range(cfg.TRAIN.EPOCH_COUNT): G.train() D.train() E.train() ZD.train() epoch_start_time = time.time() data_loader = make_dataloader(train_set, cfg.TRAIN.BATCH_SIZE, torch.cuda.current_device()) train_set.shuffle() if (epoch + 1) % 30 == 0: G_optimizer.param_groups[0]['lr'] /= 4 D_optimizer.param_groups[0]['lr'] /= 4 GE_optimizer.param_groups[0]['lr'] /= 4 ZD_optimizer.param_groups[0]['lr'] /= 4 print("learning rate change!") for y, x in data_loader: x = x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE) y_real_ = torch.ones(x.shape[0]) y_fake_ = torch.zeros(x.shape[0]) y_real_z = torch.ones( 1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0]) y_fake_z = torch.zeros( 1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0]) ############################################# D.zero_grad() D_result = D(x).squeeze() D_real_loss = BCE_loss(D_result, y_real_) z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1) z = Variable(z) x_fake = G(z).detach() D_result = D(x_fake).squeeze() D_fake_loss = BCE_loss(D_result, y_fake_) D_train_loss = D_real_loss + D_fake_loss D_train_loss.backward() D_optimizer.step() tracker.update(dict(D=D_train_loss)) ############################################# G.zero_grad() z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1) z = Variable(z) x_fake = G(z) D_result = D(x_fake).squeeze() G_train_loss = BCE_loss(D_result, y_real_) G_train_loss.backward() G_optimizer.step() tracker.update(dict(G=G_train_loss)) ############################################# ZD.zero_grad() z = torch.randn((x.shape[0], zsize)).view(-1, zsize) z = Variable(z) ZD_result = ZD(z).squeeze() ZD_real_loss = BCE_loss(ZD_result, y_real_z) z = E(x).squeeze().detach() ZD_result = ZD(z).squeeze() ZD_fake_loss = BCE_loss(ZD_result, y_fake_z) ZD_train_loss = ZD_real_loss + ZD_fake_loss ZD_train_loss.backward() ZD_optimizer.step() tracker.update(dict(ZD=ZD_train_loss)) # ############################################# E.zero_grad() G.zero_grad() z = E(x) x_d = G(z) ZD_result = ZD(z.squeeze()).squeeze() E_train_loss = BCE_loss(ZD_result, y_real_z) * 1.0 Recon_loss = F.binary_cross_entropy(x_d, x.detach()) * 2.0 (Recon_loss + E_train_loss).backward() GE_optimizer.step() tracker.update(dict(GE=Recon_loss, E=E_train_loss)) # ############################################# comparison = torch.cat([x, x_d]) save_image(comparison.cpu(), os.path.join(output_folder, 'reconstruction_' + str(epoch) + '.png'), nrow=x.shape[0]) epoch_end_time = time.time() per_epoch_ptime = epoch_end_time - epoch_start_time logger.info( '[%d/%d] - ptime: %.2f, %s' % ((epoch + 1), cfg.TRAIN.EPOCH_COUNT, per_epoch_ptime, tracker)) tracker.register_means(epoch) tracker.plot() with torch.no_grad(): resultsample = G(sample).cpu() save_image( resultsample.view(64, cfg.MODEL.INPUT_IMAGE_CHANNELS, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE), os.path.join(output_folder, 'sample_' + str(epoch) + '.png')) logger.info("Training finish!... save training results") os.makedirs("models", exist_ok=True) print("Training finish!... save training results") torch.save(G.state_dict(), "models/Gmodel_%d_%d.pkl" % (folding_id, ic)) torch.save(E.state_dict(), "models/Emodel_%d_%d.pkl" % (folding_id, ic))
model.eval() layer_count = cfg.MODEL.LAYER_COUNT decoder = nn.DataParallel(decoder) im_size = 2**(cfg.MODEL.LAYER_COUNT + 1) with torch.no_grad(): draw_uncurated_result_figure(cfg, 'make_figures/output/%s/generations.jpg' % cfg.NAME, model, cx=0, cy=0, cw=im_size, ch=im_size, rows=6, lods=[0, 0, 0, 1, 1, 2], seed=5) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-generations', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)
def encode(x): Z, _ = model.encode(x, layer_count - 1, 1) Z = Z.repeat(1, model.mapping_fl.num_layers, 1) return Z def decode(x): layer_idx = torch.arange(2 * cfg.MODEL.LAYER_COUNT)[np.newaxis, :, np.newaxis] ones = torch.ones(layer_idx.shape, dtype=torch.float32) coefs = torch.where(layer_idx < model.truncation_cutoff, 1.2 * ones, ones) x = torch.lerp(model.dlatent_avg.buff.data, x, coefs) return model.decoder(x, layer_count - 1, 1, noise=True) logger.info("Evaluating PPL metric") decoder = nn.DataParallel(decoder) with torch.no_grad(): ppl = PPL(cfg, num_samples=50000, epsilon=1e-4, space='w', sampling='full', minibatch_size=16 * torch.cuda.device_count()) ppl.evaluate(logger, mapping_fl, decoder, cfg.DATASET.MAX_RESOLUTION_LEVEL - 2, celeba_style=cfg.PPL_CELEBA_ADJUSTMENT) with torch.no_grad(): ppl = PPL(cfg, num_samples=50000, epsilon=1e-4, space='w', sampling='end', minibatch_size=16 * torch.cuda.device_count()) ppl.evaluate(logger, mapping_fl, decoder, cfg.DATASET.MAX_RESOLUTION_LEVEL - 2, celeba_style=cfg.PPL_CELEBA_ADJUSTMENT) if __name__ == "__main__": gpu_count = 1 run(sample, get_cfg_defaults(), description='ALAE-ppl', default_config='configs/ffhq.yaml', world_size=gpu_count, write_log=False)