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)
