def initialize(self, args=None, save=True):
if hasattr(args, 'name') and args.name is not None:
print("args.name= ", args.name)
self.name = args.name
if hasattr(args, 'batch_size') and args.name is not None:
self.batch_size = args.batch_size
self.main_dir = os.path.join(self.expr_dir, self.name)
self.model_dir = os.path.join(self.main_dir, "pth")
self.tb_dir = os.path.join(self.main_dir, "log")
self.info_dir = os.path.join(self.main_dir, "info")
self.output_dir = os.path.join(self.main_dir, "output")
ensure_dirs([
self.main_dir, self.model_dir, self.tb_dir, self.info_dir,
self.output_dir
])
self.device = torch.device(
"cuda:%d" % self.cuda_id if torch.cuda.is_available() else "cpu")
if save:
self.config_name = args.config
cfg_file = "%s.py" % self.config_name
shutil.copy(pjoin(BASEPATH, cfg_file),
os.path.join(self.info_dir, cfg_file))
def get_all_codes(cfg, output_path):
print(output_path)
if os.path.exists(output_path):
return np.load(output_path, allow_pickle=True)['data'].item()
ensure_dirs(os.path.dirname(output_path))
print("start over")
# Dataloader
train_loader = get_dataloader(cfg, 'train', shuffle=False)
test_loader = get_dataloader(cfg, 'test', shuffle=False)
# Trainer
trainer = Trainer(cfg)
trainer.to(cfg.device)
trainer.resume()
with torch.no_grad():
vis_dicts = {}
for phase, loader in [['train', train_loader],
['test', test_loader]]:
vis_dict = None
for t, data in enumerate(loader):
vis_codes = trainer.get_latent_codes(data)
if vis_dict is None:
vis_dict = {}
for key, value in vis_codes.items():
vis_dict[key] = [value]
else:
for key, value in vis_codes.items():
vis_dict[key].append(value)
for key, value in vis_dict.items():
if phase == "test" and key == "content_code":
continue
if key == "meta":
secondary_keys = value[0].keys()
num = len(value)
vis_dict[key] = {
secondary_key: [to_float(item) for i in range(num) for item in value[i][secondary_key]]
for secondary_key in secondary_keys}
else:
vis_dict[key] = torch.cat(vis_dict[key], 0)
vis_dict[key] = vis_dict[key].cpu().numpy()
vis_dict[key] = to_float(vis_dict[key].reshape(vis_dict[key].shape[0], -1))
vis_dicts[phase] = vis_dict
np.savez_compressed(output_path, data=vis_dicts)
return vis_dicts
def get_demo_plots(data, output_path):
"""
data: {"train": dict_train, "test": dict_test}
dict_train: {"style2d_code": blabla, etc.}
"""
ensure_dirs(output_path)
def fig_title(title):
return pjoin(output_path, title)
style_labels = data["train"]["meta"]["style"]
adain_raw = []
for key in ["style2d_adain", "style3d_adain"]:
for phase in ["train", "test"]:
adain_raw.append(data[phase][key])
adain_tsne = calc_many_blas(adain_raw, calc_tsne)
plot2D_overlay([adain_tsne[0], adain_tsne[2]],
[style_labels, style_labels], [1.0, 0.5],
fig_title(f'joint_embedding_adain_tsne'))
for key in ["style3d_code", "style3d_adain"]:
tsne_code = calc_tsne(data["train"][key])
plot2D(tsne_code, style_labels, fig_title(f'{key}_tsne'))
content_code_pca = calc_pca(data["train"]["content_code"])
indices = [
i for i in range(len(data["train"]["meta"]["content"]))
if data["train"]["meta"]["content"][i] == "walk"
]
walk_code = content_code_pca[np.array(indices)]
phase_labels = [data["train"]["meta"]["phase"][i] for i in indices]
plot2D_phase(walk_code, phase_labels, fig_title(f'content_by_phase'))
plot2D(content_code_pca, style_labels, fig_title(f'content_by_style'))
def get_all_plots(data,
output_path,
writers,
iter,
summary=True,
style_cluster_protocols=('pca'),
separate_compute=False):
"""
data: {"train": dict_train, "test": dict_test}
dict_train: {"style2d_code": blabla, etc.}
separate_compute: compute t-SNE for 2D & 3D separately
"""
ensure_dirs(output_path)
def fig_title(title):
return pjoin(output_path, title)
def add_fig(fig, title, phase):
if summary:
writers[phase].add_figure(title, fig, global_step=iter)
keys = data["train"].keys()
has2d = "style2d_code" in keys
has3d = "style3d_code" in keys
# style codes & adain params
for suffix in ["_code", "_adain"]:
codes_raw = []
titles = []
phases = []
data_keys = []
if has2d: data_keys.append("style2d" + suffix)
if has3d: data_keys.append("style3d" + suffix)
for key in data_keys:
for phase in ["train", "test"]:
codes_raw.append(data[phase][key])
titles.append(f'{phase}_{key}')
phases.append(phase)
# calc tsne with style2/3d, train/test altogether
for name, protocol in zip(['pca', 'tsne'], [calc_pca, calc_tsne]):
if name not in style_cluster_protocols:
continue
style_codes = calc_many_blas(codes_raw, protocol)
fig = plot2D_overlay([style_codes[0], style_codes[2]], [
data["train"]["meta"]["style"], data["train"]["meta"]["style"]
], [1.0, 0.5], fig_title(f'joint_embedding_{name}{suffix}'))
add_fig(fig, f'joint_embedding_{name}{suffix}', "train")
for i, (code, phase,
title) in enumerate(zip(style_codes, phases, titles)):
if separate_compute:
code = protocol(codes_raw[i])
for label_type in ["style", "content"]:
fig = plot2D(code, data[phase]["meta"][label_type],
fig_title(f'{title}_{name}_{label_type}'))
add_fig(fig, f'{title}_{name}_{label_type}', phase)
# content codes (train only)
content_code_pca = calc_pca(data["train"]["content_code"])
for label in ["style", "content", "phase"]:
if label == "phase":
indices = [
i for i in range(len(data["train"]["meta"]["content"]))
if data["train"]["meta"]["content"][i] == "walk"
]
walk_code = content_code_pca[np.array(indices)]
phase_labels = [data["train"]["meta"]["phase"][i] for i in indices]
fig = plot2D_phase(walk_code, phase_labels,
fig_title(f'content_by_{label}'))
else:
fig = plot2D(content_code_pca, data["train"]["meta"][label],
fig_title(f'content_by_{label}'))
add_fig(fig, f'content_by_{label}', "train")
"""