def test_result_subdir(self):
result_dir = '/tmp'
result_subdir = util.create_result_subdir(result_dir,
desc='test',
profile={})
util.locate_result_subdir(result_dir, result_subdir)
util.locate_result_subdir(result_dir, 0)
util.locate_result_subdir(result_dir, '000')
def sample(ctx):
hps = ctx.obj['hps']
inferer = ctx.obj['inferer']
# smaple
img = inferer.sample(z=None, y_onehot=None, eps_std=0.5)
# save result
result_subdir = util.create_result_subdir(hps.general.result_dir,
desc='sample',
profile=hps)
util.tensor_to_pil(img).save(os.path.join(result_subdir, 'sample.png'))
def compute_deltaz(ctx):
hps = ctx.obj['hps']
inferer = ctx.obj['inferer']
dataset = ctx.obj['dataset']
# compute delta
deltaz = inferer.compute_attribute_delta(dataset)
# save result
result_subdir = util.create_result_subdir(hps.general.result_dir,
desc='deltaz',
profile=hps)
util.save_deltaz(deltaz, result_subdir)
def interpolate(ctx, delta_file, image_file, batch):
hps = ctx.obj['hps']
inferer = ctx.obj['inferer']
dataset = ctx.obj['dataset']
img = Image.open(image_file).convert('RGB')
deltaz = util.load_deltaz(delta_file)
result_subdir = util.create_result_subdir(hps.general.result_dir,
desc='interpolation',
profile=hps)
if batch:
interpolation_vector = util.make_interpolation_vector(
hps.dataset.num_classes)
for cls in range(interpolation_vector.shape[0]):
print('[Inferer] interpolating class "{}"'.format(
dataset.attrs[cls]))
imgs_interpolated = []
progress = tqdm(range(interpolation_vector.shape[1]))
for lv in progress:
img_interpolated = inferer.apply_attribute_delta(
img, deltaz, interpolation_vector[cls, lv, :])
imgs_interpolated.append(img_interpolated)
# img_interpolated = util.tensor_to_pil(img_interpolated)
# img_interpolated.save('interpolation/interpolated_{:s}_{:0.2f}.png'.format(
# dataset.attrs[cls],
# interpolation_vector[cls, lv, cls]))
imgs_stacked = torch.stack(imgs_interpolated)
imgs_grid = make_grid(imgs_stacked,
nrow=interpolation_vector.shape[1])
imgs = util.tensor_to_pil(imgs_grid)
imgs.save(
os.path.join(
result_subdir,
'interpolated_{:s}.png'.format(dataset.attrs[cls])))
else:
interpolation = [0.] * hps.dataset.num_classes
interpolation[0] = 1.
img_interpolated = inferer.apply_attribute_delta(
img, deltaz, interpolation)
img_interpolated = util.tensor_to_pil(img_interpolated)
img_interpolated.save(os.path.join(result_subdir, 'interpolated.png'))
def reconstruct(ctx, image_path):
hps = ctx.obj['hps']
inferer = ctx.obj['inferer']
# image_path = "/home/nird/glow_pytorch/experiments_images/"
# get image list
img_list = []
if os.path.isfile(image_path) and util.is_image(image_path):
img_list = [image_path]
elif os.path.isdir(image_path):
img_list = [
os.path.join(image_path, f) for f in os.listdir(image_path)
if util.is_image(os.path.join(image_path, f))
]
# reconstruct images
img_grid_list = []
util.check_path('reconstructed')
for img_path in img_list:
img = Image.open(img_path).convert('RGB')
x = util.pil_to_tensor(img,
transform=transforms.Compose(
(transforms.CenterCrop(160),
transforms.Resize(128),
transforms.ToTensor())))
z = inferer.encode(x)
x_ = inferer.decode(z)
img_grid = torch.cat((x, x_.cpu()), dim=1)
img_grid_list.append(img_grid)
# util.tensor_to_pil(img_grid).save('reconstructed/{}'.format(os.path.basename(img_path)))
# generate grid of reconstructed images
imgs_grid = make_grid(torch.stack(img_grid_list))
# save result
result_subdir = util.create_result_subdir(hps.general.result_dir,
desc='reconstruct',
profile=hps)
util.tensor_to_pil(imgs_grid).save(os.path.join(result_subdir, 'grid.png'))
def build(self, training=True):
"""
Build network
:param training:
:type training:
:return:
:rtype:
"""
# initialize all variables
step = 0
state = None
result_subdir = None
graph, optimizer, scheduler, criterion_dict = None, None, None, None
devices = util.get_devices(self.hps.device.graph)
data_device = util.get_devices(self.hps.device.data)[0]
# build graph
graph = model.Glow(self.hps)
graph.to('cpu')
# load model
if graph is not None:
# locate or create result subdir
if self.hps.general.warm_start and self.hps.general.resume_run_id != "":
result_subdir = util.locate_result_subdir(
self.hps.general.result_dir,
self.hps.general.resume_run_id)
if training and result_subdir is None:
result_subdir = util.create_result_subdir(
self.hps.general.result_dir,
desc=self.hps.profile,
profile=self.hps)
# load pre-trained model on first device
if self.hps.general.warm_start:
step_or_model_path = None
if os.path.exists(self.hps.general.pre_trained):
step_or_model_path = self.hps.general.pre_trained
elif self.hps.general.resume_step not in [
'', 'best', 'latest'
]:
step_or_model_path = int(self.hps.general.resume_step)
if step_or_model_path is not None:
state = util.load_model(result_subdir,
step_or_model_path,
graph,
device=devices[0])
if not training and state is None:
raise RuntimeError('No pre-trained model for inference')
# move graph to devices
if 'cpu' in devices:
graph = graph.cpu()
data_device = 'cpu'
else:
graph = graph.to(devices[0])
print('[Builder] Use {} for model running and {} for data loading'.
format(devices[0], data_device))
# setup optimizer and lr scheduler
if training and graph is not None:
# get optimizer
optimizer_name = self.hps.optim.optimizer.lower()
assert optimizer_name in self.optimizer_dict.keys(), \
"Unsupported optimizer: {}".format(optimizer_name)
# If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it.
optimizer = self.optimizer_dict[optimizer_name](
graph.parameters(), **self.hps.optim.optimizer_args)
if state is not None:
optimizer.load_state_dict(state['optimizer'])
# get lr scheduler
scheduler_name = self.hps.optim.lr_scheduler.lower()
scheduler_args = self.hps.optim.lr_scheduler_args
assert scheduler_name in self.lr_scheduler_dict.keys(), \
"Unsupported lr scheduler: {}".format(scheduler_name)
if 'base_lr' not in scheduler_args:
scheduler_args['base_lr'] = self.hps.optim.optimizer_args['lr']
scheduler = partial(self.lr_scheduler_dict[scheduler_name],
**scheduler_args)
return {
'step': step,
'graph': graph,
'optimizer': optimizer,
'scheduler': scheduler,
'devices': devices,
'data_device': data_device,
'result_subdir': result_subdir
}
