def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
z = xp.random.normal(0,
args.temperature,
size=(
1,
3,
) + hyperparams.image_size).astype("float32")
x, _ = decoder.reverse_step(z)
x_img = make_uint8(x.data[0], num_bins_x)
plt.imshow(x_img, interpolation="none")
plt.pause(.01)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
image_size = (28, 28)
images = chainer.datasets.mnist.get_mnist(withlabel=False)[0]
images = 255.0 * np.asarray(images).reshape((-1, ) + image_size + (1, ))
if hyperparams.num_image_channels != 1:
images = np.broadcast_to(images, (images.shape[0], ) + image_size +
(hyperparams.num_image_channels, ))
images = preprocess(images, hyperparams.num_bits_x)
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=1)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
fig = plt.figure(figsize=(8, 4))
left = fig.add_subplot(1, 2, 1)
right = fig.add_subplot(1, 2, 2)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
for data_indices in iterator:
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
factorized_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
factorized_z.append(zi)
# for zi in factorized_z:
# noise = xp.random.normal(
# 0, 0.2, size=zi.shape).astype("float32")
# zi.data += noise
rev_x, _ = decoder.reverse_step(factorized_z)
x_img = make_uint8(x[0], num_bins_x)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
left.imshow(x_img, interpolation="none")
right.imshow(rev_x_img, interpolation="none")
plt.pause(.01)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, hyperparams.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, hyperparams.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=1)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
for data_indices in iterator:
print("data:", data_indices)
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
for (_, mean, ln_var) in factorized_z_distribution:
print(xp.mean(mean.data), xp.mean(xp.exp(ln_var.data)))
def get_model(path, using_gpu):
print(path)
hyperparams = Hyperparameters(path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=path)
if using_gpu:
encoder.to_gpu()
return encoder, num_bins_x, hyperparams
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
temperatures = [0.0, 0.25, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
total = len(temperatures)
fig = plt.figure(figsize=(total * 4, 4))
subplots = []
for n in range(total):
subplot = fig.add_subplot(1, total, n + 1)
subplots.append(subplot)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
z_batch = []
for temperature in temperatures:
z = np.random.normal(0,
temperature,
size=(3, ) +
hyperparams.image_size).astype("float32")
z_batch.append(z)
z_batch = np.asanyarray(z_batch)
if using_gpu:
z_batch = cuda.to_gpu(z_batch)
x, _ = decoder.reverse_step(z_batch)
for n, (temperature,
subplot) in enumerate(zip(temperatures, subplots)):
x_img = make_uint8(x.data[n], num_bins_x)
# x_img = np.broadcast_to(x_img, (28, 28, 3))
subplot.imshow(x_img, interpolation="none")
subplot.set_title("temperature={}".format(temperature))
plt.pause(.01)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
total = hyperparams.levels
fig = plt.figure(figsize=(4 * total, 4))
subplots = []
for n in range(total):
subplot = fig.add_subplot(1, total, n + 1)
subplots.append(subplot)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
seed = int(time.time())
for level in range(1, hyperparams.levels):
xp.random.seed(seed)
z = xp.random.normal(0,
args.temperature,
size=(
1,
3,
) + hyperparams.image_size,
dtype="float32")
factorized_z = glow.nn.functions.factor_z(
z, level + 1, squeeze_factor=hyperparams.squeeze_factor)
out = glow.nn.functions.unsqueeze(
factorized_z.pop(-1),
factor=hyperparams.squeeze_factor,
module=xp)
for n, zi in enumerate(factorized_z[::-1]):
block = encoder.blocks[level - n - 1]
out, _ = block.reverse_step(
out,
gaussian_eps=zi,
squeeze_factor=hyperparams.squeeze_factor)
rev_x = out
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplot = subplots[level - 1]
subplot.imshow(rev_x_img, interpolation="none")
subplot.set_title("level = {}".format(level))
# original #levels
xp.random.seed(seed)
z = xp.random.normal(0,
args.temperature,
size=(
1,
3,
) + hyperparams.image_size,
dtype="float32")
factorized_z = encoder.factor_z(z)
rev_x, _ = decoder.reverse_step(factorized_z)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplot = subplots[-1]
subplot.imshow(rev_x_img, interpolation="none")
subplot.set_title("level = {}".format(hyperparams.levels))
plt.pause(.01)
def main():
try:
os.mkdir(args.snapshot_path)
except:
pass
comm = chainermn.create_communicator()
device = comm.intra_rank
print("device", device, "/", comm.size)
cuda.get_device(device).use()
xp = cupy
num_bins_x = 2**args.num_bits_x
images = None
if comm.rank == 0:
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(
args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, args.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, args.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
assert args.image_size == images.shape[2]
x_mean = np.mean(images)
x_var = np.var(images)
print(
tabulate([
["#", len(images)],
["mean", x_mean],
["var", x_var],
]))
dataset = chainermn.scatter_dataset(images, comm, shuffle=True)
hyperparams = Hyperparameters()
hyperparams.levels = args.levels
hyperparams.depth_per_level = args.depth_per_level
hyperparams.nn_hidden_channels = args.nn_hidden_channels
hyperparams.image_size = (args.image_size, args.image_size)
hyperparams.num_bits_x = args.num_bits_x
hyperparams.lu_decomposition = args.lu_decomposition
hyperparams.save(args.snapshot_path)
if comm.rank == 0:
hyperparams.save(args.snapshot_path)
hyperparams.print()
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
encoder.to_gpu()
optimizer = chainermn.create_multi_node_optimizer(
chainer.optimizers.Adam(alpha=1e-4), comm)
optimizer.setup(encoder)
current_training_step = 0
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
# Training loop
for iteration in range(args.total_iteration):
sum_loss = 0
sum_nll = 0
total_batch = 0
start_time = time.time()
iterator = chainer.iterators.SerialIterator(
dataset, args.batch_size, repeat=False)
# Data dependent initialization
if encoder.need_initialize:
for data in iterator:
x = to_gpu(np.asanyarray(data))
encoder.initialize_actnorm_weights(x)
break
for batch_index, data in enumerate(iterator):
x = to_gpu(np.asanyarray(data))
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
batch_size = x.shape[0]
denom = math.log(2.0) * num_pixels
factorized_z_distribution, logdet = encoder.forward_step(x)
logdet -= math.log(num_bins_x) * num_pixels
negative_log_likelihood = 0
for (zi, mean, ln_var) in factorized_z_distribution:
negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var)
loss = (negative_log_likelihood / batch_size - logdet) / denom
encoder.cleargrads()
loss.backward()
optimizer.update()
current_training_step += 1
total_batch += 1
sum_loss += float(loss.data)
sum_nll += float(negative_log_likelihood.data) / args.batch_size
if comm.rank == 0:
printr(
"Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - log_det: {:.8f}".
format(
iteration + 1, batch_index + 1,
len(dataset) // batch_size, float(loss.data),
float(negative_log_likelihood.data) / batch_size /
denom,
float(logdet.data) / denom))
if (batch_index + 1) % 100 == 0:
encoder.save(args.snapshot_path)
if comm.rank == 0:
log_likelihood = -sum_nll / total_batch
elapsed_time = time.time() - start_time
print(
"\033[2KIteration {} - loss: {:.5f} - log_likelihood: {:.5f} - elapsed_time: {:.3f} min".
format(iteration + 1, sum_loss / total_batch, log_likelihood,
elapsed_time / 60))
encoder.save(args.snapshot_path)
def main():
try:
os.mkdir(args.ckpt)
except:
pass
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
# Load picture
x = np.array(Image.open(args.img)).astype('float32')
x = preprocess(x, hyperparams.num_bits_x)
x = to_gpu(xp.expand_dims(x, axis=0))
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
if True:
# Print this image info:
b = xp.zeros((1, 3, 128, 128))
z, fw_ldt = encoder.forward_step(x, b)
fw_ldt -= math.log(num_bins_x) * num_pixels
logpZ = 0
ez = []
factor_z = []
for (zi, mean, ln_var) in z:
factor_z.append(zi.data)
logpZ += cf.gaussian_nll(zi, mean, ln_var)
ez.append(zi.data.reshape(-1, ))
ez = np.concatenate(ez)
logpZ2 = cf.gaussian_nll(ez, xp.zeros(ez.shape),
xp.zeros(ez.shape)).data
print(fw_ldt, logpZ, logpZ2)
with encoder.reverse() as decoder:
rx, _ = decoder.reverse_step(factor_z)
rx_img = make_uint8(rx.data[0], num_bins_x)
rx_img = Image.fromarray(rx_img)
rx_img.save(args.t + 'ori_revx.png')
np.save(args.t + 'ori_z.npy', ez.get())
# Construct epsilon
class eps(chainer.Chain):
def __init__(self, shape, glow_encoder):
super().__init__()
self.encoder = glow_encoder
with self.init_scope():
self.b = chainer.Parameter(initializers.Zero(),
(1, 3, 128, 128))
self.m = chainer.Parameter(initializers.One(), (3, 8, 8))
def forward(self, x):
# b = cf.tanh(self.b) * 0.5
b = self.b
# Not sure if implementation is wrong
m = cf.softplus(self.m)
# m = cf.repeat(m, 8, axis=2)
# m = cf.repeat(m, 8, axis=1)
# m = cf.repeat(m, 16, axis=2)
# m = cf.repeat(m, 16, axis=1)
# b = b * m
# cur_x = cf.add(x, b)
# cur_x = cf.clip(cur_x, -0.5,0.5)
z = []
zs, logdet = self.encoder.forward_step(x, b)
for (zi, mean, ln_var) in zs:
z.append(zi)
z = merge_factorized_z(z)
# return z, zs, logdet, cf.batch_l2_norm_squared(b), xp.tanh(self.b.data*1), cur_x, m
return z, zs, logdet, xp.sum(xp.abs(b.data)), self.b.data * 1, m, x
def save(self, path):
filename = 'loss_model.hdf5'
self.save_parameter(path, filename, self)
def save_parameter(self, path, filename, params):
tmp_filename = str(uuid.uuid4())
tmp_filepath = os.path.join(path, tmp_filename)
save_hdf5(tmp_filepath, params)
os.rename(tmp_filepath, os.path.join(path, filename))
epsilon = eps(x.shape, encoder)
if using_gpu:
epsilon.to_gpu()
# optimizer = Optimizer(epsilon)
optimizer = optimizers.Adam().setup(epsilon)
# optimizer = optimizers.SGD().setup(epsilon)
epsilon.b.update_rule.hyperparam.lr = 0.01
epsilon.m.update_rule.hyperparam.lr = 0.1
print('init finish')
training_step = 0
z_s = []
b_s = []
loss_s = []
logpZ_s = []
logDet_s = []
m_s = []
j = 0
for iteration in range(args.total_iteration):
epsilon.cleargrads()
z, zs, fw_ldt, b_norm, b, m, cur_x = epsilon.forward(x)
fw_ldt -= math.log(num_bins_x) * num_pixels
logpZ1 = 0
factor_z = []
for (zi, mean, ln_var) in zs:
factor_z.append(zi.data)
logpZ1 += cf.gaussian_nll(zi, mean, ln_var)
logpZ2 = cf.gaussian_nll(z, xp.zeros(z.shape), xp.zeros(z.shape)).data
# logpZ2 = cf.gaussian_nll(z, np.mean(z), np.log(np.var(z))).data
logpZ = (logpZ2 + logpZ1) / 2
loss = b_norm + (logpZ - fw_ldt)
loss.backward()
optimizer.update()
training_step += 1
z_s.append(z.get())
b_s.append(cupy.asnumpy(b))
m_s.append(cupy.asnumpy(m.data))
loss_s.append(_float(loss))
logpZ_s.append(_float(logpZ))
logDet_s.append(_float(fw_ldt))
printr(
"Iteration {}: loss: {:.6f} - b_norm: {:.6f} - logpZ: {:.6f} - logpZ1: {:.6f} - logpZ2: {:.6f} - log_det: {:.6f} - logpX: {:.6f}\n"
.format(iteration + 1, _float(loss), _float(b_norm), _float(logpZ),
_float(logpZ1), _float(logpZ2), _float(fw_ldt),
_float(logpZ) - _float(fw_ldt)))
if iteration % 100 == 99:
np.save(args.ckpt + '/' + str(j) + 'z.npy', z_s)
np.save(args.ckpt + '/' + str(j) + 'b.npy', b_s)
np.save(args.ckpt + '/' + str(j) + 'loss.npy', loss_s)
np.save(args.ckpt + '/' + str(j) + 'logpZ.npy', logpZ_s)
np.save(args.ckpt + '/' + str(j) + 'logDet.npy', logDet_s)
# cur_x = make_uint8(cur_x[0].data, num_bins_x)
# np.save(args.ckpt + '/'+str(j)+'image.npy', cur_x)
np.save(args.ckpt + '/' + str(j) + 'm.npy', m_s)
with encoder.reverse() as decoder:
rx, _ = decoder.reverse_step(factor_z)
rx_img = make_uint8(rx.data[0], num_bins_x)
np.save(args.ckpt + '/' + str(j) + 'res.npy', rx_img)
z_s = []
b_s = []
loss_s = []
logpZ_s = []
logDet_s = []
m_s = []
j += 1
epsilon.save(args.ckpt)
def main():
try:
os.mkdir(args.snapshot_path)
except:
pass
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
num_bins_x = 2**args.num_bits_x
image_size = (28, 28)
images = chainer.datasets.mnist.get_mnist(withlabel=False)[0]
images = 255.0 * np.asarray(images).reshape((-1, ) + image_size + (1, ))
if args.num_channels != 1:
images = np.broadcast_to(
images, (images.shape[0], ) + image_size + (args.num_channels, ))
images = preprocess(images, args.num_bits_x)
x_mean = np.mean(images)
x_var = np.var(images)
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=args.batch_size)
print(tabulate([
["#", len(dataset)],
["mean", x_mean],
["var", x_var],
]))
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.levels = args.levels
hyperparams.depth_per_level = args.depth_per_level
hyperparams.nn_hidden_channels = args.nn_hidden_channels
hyperparams.image_size = image_size
hyperparams.num_bits_x = args.num_bits_x
hyperparams.lu_decomposition = args.lu_decomposition
hyperparams.num_image_channels = args.num_channels
hyperparams.save(args.snapshot_path)
hyperparams.print()
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
optimizer = Optimizer(encoder)
# Data dependent initialization
if encoder.need_initialize:
for batch_index, data_indices in enumerate(iterator):
x = to_gpu(dataset[data_indices])
encoder.initialize_actnorm_weights(
x, reduce_memory=args.reduce_memory)
break
current_training_step = 0
num_pixels = args.num_channels * hyperparams.image_size[0] * hyperparams.image_size[1]
# Training loop
for iteration in range(args.total_iteration):
sum_loss = 0
sum_nll = 0
start_time = time.time()
for batch_index, data_indices in enumerate(iterator):
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
denom = math.log(2.0) * num_pixels
factorized_z_distribution, logdet = encoder.forward_step(
x, reduce_memory=args.reduce_memory)
logdet -= math.log(num_bins_x) * num_pixels
negative_log_likelihood = 0
for (zi, mean, ln_var) in factorized_z_distribution:
negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var)
loss = (negative_log_likelihood / args.batch_size - logdet) / denom
encoder.cleargrads()
loss.backward()
optimizer.update(current_training_step)
current_training_step += 1
sum_loss += float(loss.data)
sum_nll += float(negative_log_likelihood.data) / args.batch_size
printr(
"Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - log_det: {:.8f}".
format(
iteration + 1, batch_index + 1, len(iterator),
float(loss.data),
float(negative_log_likelihood.data) / args.batch_size /
denom,
float(logdet.data) / denom))
log_likelihood = -sum_nll / len(iterator)
elapsed_time = time.time() - start_time
print(
"\033[2KIteration {} - loss: {:.5f} - log_likelihood: {:.5f} - elapsed_time: {:.3f} min".
format(iteration + 1, sum_loss / len(iterator), log_likelihood,
elapsed_time / 60))
encoder.save(args.snapshot_path)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
image_size = (28, 28)
images = chainer.datasets.mnist.get_mnist(withlabel=False)[0]
images = 255.0 * np.asarray(images).reshape((-1, ) + image_size + (1, ))
if hyperparams.num_image_channels != 1:
images = np.broadcast_to(images, (images.shape[0], ) + image_size +
(hyperparams.num_image_channels, ))
images = preprocess(images, hyperparams.num_bits_x)
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=2)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
total = args.num_steps + 2
fig = plt.figure(figsize=(4 * total, 4))
subplots = []
for n in range(total):
subplot = fig.add_subplot(1, total, n + 1)
subplots.append(subplot)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
for data_indices in iterator:
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
factorized_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
factorized_z.append(zi)
z = encoder.merge_factorized_z(factorized_z)
z_start = z[0]
z_end = z[1]
z_batch = [z_start]
for n in range(args.num_steps):
ratio = n / (args.num_steps - 1)
z_interp = ratio * z_end + (1.0 - ratio) * z_start
z_batch.append(args.temperature * z_interp)
z_batch.append(z_end)
z_batch = xp.stack(z_batch)
rev_x_batch, _ = decoder.reverse_step(z_batch)
for n in range(args.num_steps):
rev_x_img = make_uint8(rev_x_batch.data[n + 1], num_bins_x)
subplots[n + 1].imshow(rev_x_img, interpolation="none")
x_start_img = make_uint8(x[0], num_bins_x)
subplots[0].imshow(x_start_img, interpolation="none")
x_end_img = make_uint8(x[-1], num_bins_x)
subplots[-1].imshow(x_end_img, interpolation="none")
plt.pause(.01)
def main():
try:
os.mkdir(args.snapshot_path)
except:
pass
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
num_bins_x = 2**args.num_bits_x
assert args.dataset_format in ["png", "npy"]
# Get datasets:
if True:
files = Path(args.dataset_path).glob("*.{}".format(
args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, args.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
# TODO: Preprocess here
array = preprocess(array, args.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
# Print dataset information
if True:
x_mean = np.mean(images)
x_var = np.var(images)
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=args.batch_size)
print(
tabulate([
["#", len(dataset)],
["mean", x_mean],
["var", x_var],
]))
# Hyperparameters' info
if True:
hyperparams = Hyperparameters()
hyperparams.levels = args.levels
hyperparams.depth_per_level = args.depth_per_level
hyperparams.nn_hidden_channels = args.nn_hidden_channels
hyperparams.image_size = images.shape[2:]
hyperparams.num_bits_x = args.num_bits_x
hyperparams.lu_decomposition = args.lu_decomposition
hyperparams.squeeze_factor = args.squeeze_factor
hyperparams.save(args.snapshot_path)
hyperparams.print()
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
optimizer = Optimizer(encoder)
# Data dependent initialization
if encoder.need_initialize:
for batch_index, data_indices in enumerate(iterator):
x = to_gpu(dataset[data_indices])
encoder.initialize_actnorm_weights(x)
break
current_training_step = 0
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
# Training loop
for iteration in range(args.total_iteration):
sum_loss = 0
sum_nll = 0
sum_kld = 0
start_time = time.time()
for batch_index, data_indices in enumerate(iterator):
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
denom = math.log(2.0) * num_pixels
factorized_z_distribution, logdet = encoder.forward_step(x)
logdet -= math.log(num_bins_x) * num_pixels
kld = 0
negative_log_likelihood = 0
factor_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var)
if args.regularize_z:
kld += cf.gaussian_kl_divergence(mean, ln_var)
factor_z.append(zi.data.reshape(zi.shape[0], -1))
factor_z = xp.concatenate(factor_z, axis=1)
negative_log_likelihood += cf.gaussian_nll(
factor_z, xp.zeros(factor_z.shape, dtype='float32'),
xp.zeros(factor_z.shape, dtype='float32'))
loss = (negative_log_likelihood + kld) / args.batch_size - logdet
loss = loss / denom
encoder.cleargrads()
loss.backward()
optimizer.update(current_training_step)
current_training_step += 1
sum_loss += _float(loss)
sum_nll += _float(negative_log_likelihood) / args.batch_size
sum_kld += _float(kld) / args.batch_size
printr(
"Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - kld: {:.8f} - log_det: {:.8f}\n"
.format(
iteration + 1, batch_index + 1, len(iterator),
_float(loss),
_float(negative_log_likelihood) / args.batch_size / denom,
_float(kld) / args.batch_size,
_float(logdet) / denom))
if (batch_index + 1) % 100 == 0:
encoder.save(args.snapshot_path)
mean_log_likelihood = -sum_nll / len(iterator)
mean_kld = sum_kld / len(iterator)
elapsed_time = time.time() - start_time
print(
"\033[2KIteration {} - loss: {:.5f} - log_likelihood: {:.5f} - kld: {:.5f} - elapsed_time: {:.3f} min\n"
.format(iteration + 1, sum_loss / len(iterator),
mean_log_likelihood, mean_kld, elapsed_time / 60))
encoder.save(args.snapshot_path)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
# Get Dataset:
if True:
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(
args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, hyperparams.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, hyperparams.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=1)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
ori_x = []
enc_z = []
rev_x = []
fw_logdet = []
logpZ = []
logpZ2 = []
i = 0
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
for data_indices in iterator:
i += 1
x = to_gpu(dataset[data_indices]) # 1x3x64x64
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
x_img = make_uint8(x[0], num_bins_x)
ori_x.append(x_img)
factorized_z_distribution, fw_ldt = encoder.forward_step(x)
fw_ldt -= math.log(num_bins_x) * num_pixels
fw_logdet.append(cupy.asnumpy(fw_ldt.data))
factor_z = []
ez = []
nll = 0
for (zi, mean, ln_var) in factorized_z_distribution:
nll += cf.gaussian_nll(zi, mean, ln_var)
factor_z.append(zi.data)
ez.append(zi.data.reshape(-1, ))
ez = np.concatenate(ez)
enc_z.append(ez.get())
logpZ.append(cupy.asnumpy(nll.data))
logpZ2.append(
cupy.asnumpy(
cf.gaussian_nll(ez, np.mean(ez), np.log(np.var(ez))).data))
rx, _ = decoder.reverse_step(factor_z)
rx_img = make_uint8(rx.data[0], num_bins_x)
rev_x.append(rx_img)
if i % 100 == 0:
np.save(str(i) + '/ori_x.npy', ori_x)
fw_logdet = np.array(fw_logdet)
np.save(str(i) + '/fw_logdet.npy', fw_logdet)
np.save(str(i) + '/enc_z.npy', enc_z)
logpZ = np.array(logpZ)
np.save(str(i) + '/logpZ.npy', logpZ)
logpZ2 = np.array(logpZ2)
np.save(str(i) + '/logpZ2.npy', logpZ2)
np.save(str(i) + '/rev_x.npy', rev_x)
ori_x = []
enc_z = []
rev_x = []
fw_logdet = []
logpZ = []
logpZ2 = []
return
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
# while True:
# z = xp.random.normal(
# 0, args.temperature, size=(
# 1,
# 3,
# ) + hyperparams.image_size).astype("float32")
# x, _ = decoder.reverse_step(z)
# x_img = make_uint8(x.data[0], num_bins_x)
# plt.imshow(x_img, interpolation="none")
# plt.pause(.01)
i = 0
j = 0
enc_z = []
rev_x = []
bk_logdet = []
logpZ2 = []
fw_logdet = []
sec_z = []
sec_pz = []
sec_pz2 = []
while j < 6:
i += 1
z = xp.random.normal(0,
args.temperature,
size=(
1,
3,
) + hyperparams.image_size).astype("float32")
enc_z.append(cupy.asnumpy(z))
lvar = xp.log(args.temperature)
logpZ2.append(cupy.asnumpy(cf.gaussian_nll(z, 0, lvar).data))
x, blogd = decoder.reverse_step(z)
x_img = make_uint8(x.data[0], num_bins_x)
rev_x.append(x_img)
bk_logdet.append(cupy.asnumpy(blogd.data))
factorized_z_distribution, fw_ldt = encoder.forward_step(x)
fw_logdet.append(cupy.asnumpy(fw_ldt.data))
factor_z = []
ez = []
nll = 0
for (zi, mean, ln_var) in factorized_z_distribution:
nll += cf.gaussian_nll(zi, mean, ln_var)
factor_z.append(zi.data)
ez.append(zi.data.reshape(-1, ))
ez = np.concatenate(ez)
sec_z.append(ez.get())
sec_pz.append(cupy.asnumpy(nll.data))
sec_pz2.append(
cupy.asnumpy(
cf.gaussian_nll(ez, np.mean(ez), np.log(np.var(ez))).data))
if i % 250 == 0:
i = 0
j += 1
print('dataset: ', j)
np.save('sample/' + str(j) + 'enc_z.npy', enc_z)
np.save('sample/' + str(j) + 'rev_x.npy', rev_x)
bk_logdet = cupy.asnumpy(bk_logdet)
np.save('sample/' + str(j) + 'bk_logdet.npy', bk_logdet)
logpZ2 = cupy.asnumpy(logpZ2)
np.save('sample/' + str(j) + 'logpZ2.npy', logpZ2)
fw_logdet = cupy.asnumpy(fw_logdet)
np.save('sample/' + str(j) + 'fw_logdet.npy', fw_logdet)
np.save('sample/' + str(j) + 'sec_z.npy', sec_z)
sec_pz = cupy.asnumpy(sec_pz)
np.save('sample/' + str(j) + 'sec_pz.npy', sec_pz)
sec_pz2 = cupy.asnumpy(sec_pz2)
np.save('sample/' + str(j) + 'sec_pz2.npy', sec_pz2)
enc_z = []
rev_x = []
bk_logdet = []
logpZ2 = []
fw_logdet = []
sec_z = []
sec_pz = []
sec_pz2 = []
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, hyperparams.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, hyperparams.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=1)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
fig = plt.figure(figsize=(8, 4))
left = fig.add_subplot(1, 2, 1)
right = fig.add_subplot(1, 2, 2)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
for data_indices in iterator:
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
factorized_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
factorized_z.append(zi)
rev_x, _ = decoder.reverse_step(factorized_z)
x_img = make_uint8(x[0], num_bins_x)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
left.imshow(x_img, interpolation="none")
right.imshow(rev_x_img, interpolation="none")
plt.pause(.01)
def main():
try:
os.mkdir(args.ckpt)
except:
pass
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
# Load picture
x = np.array(Image.open(args.img)).astype('float32')
x = preprocess(x, hyperparams.num_bits_x)
x = to_gpu(xp.expand_dims(x, axis=0))
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
# Construct epsilon
class eps(chainer.Chain):
def __init__(self, shape, glow_encoder):
super().__init__()
self.encoder = glow_encoder
with self.init_scope():
self.b = chainer.Parameter(initializers.Zero(), shape)
self.m = chainer.Parameter(initializers.One(), shape)
def modify_mask(self):
mask = self.m.data
for i_idx in range(8):
for j_idx in range(8):
mean = xp.mean((xp.sum(mask[:, :, i_idx * 8:i_idx * 8 + 8,
j_idx * 8:j_idx * 8 + 8])))
mask[:, :, i_idx * 8:i_idx * 8 + 8,
j_idx * 8:j_idx * 8 + 8] = mean
mask = xp.abs(mask)
print(type(mask), type(self.m), type(self.m.data))
self.m.data = mask
def forward(self, x):
# b_ = cf.tanh(self.b)
b_ = self.b
# Not sure if implementation is wrong
self.modify_mask()
# m = cf.repeat(m, 8, axis=2)
# m = cf.repeat(m, 8, axis=1)
# m = cf.repeat(m, 16, axis=2)
# m = cf.repeat(m, 16, axis=1)
# b = b * m
x_ = cf.add(x, b_)
x_ = cf.clip(x_, -0.5, 0.5)
z = []
zs, logdet = self.encoder.forward_step(x_)
for (zi, mean, ln_var) in zs:
z.append(zi)
z = merge_factorized_z(z)
# return z, zs, logdet, cf.batch_l2_norm_squared(b), xp.tanh(self.b.data*1), cur_x, m
return z, zs, logdet, xp.sum(xp.abs(b_.data)), xp.tanh(
self.b.data * 1), self.m, x_
def save(self, path):
filename = 'loss_model.hdf5'
self.save_parameter(path, filename, self)
def save_parameter(self, path, filename, params):
tmp_filename = str(uuid.uuid4())
tmp_filepath = os.path.join(path, tmp_filename)
save_hdf5(tmp_filepath, params)
os.rename(tmp_filepath, os.path.join(path, filename))
epsilon = eps(x.shape, encoder)
if using_gpu:
epsilon.to_gpu()
# optimizer = Optimizer(epsilon)
# optimizer = optimizers.Adam(alpha=0.0005).setup(epsilon)
optimizer = optimizers.SGD().setup(epsilon)
epsilon.b.update_rule.hyperparam.lr = 0.0001
epsilon.m.update_rule.hyperparam.lr = 0.1
print('init finish')
training_step = 0
z_s = []
b_s = []
loss_s = []
logpZ_s = []
logDet_s = []
m_s = []
j = 0
for iteration in range(args.total_iteration):
# z, zs, logdet, xp.sum(xp.abs(b_.data)), xp.tanh(self.b.data * 1), self.m, x_
z, zs, fw_ldt, b_norm, b, m, cur_x = epsilon.forward(x)
epsilon.cleargrads()
fw_ldt -= math.log(num_bins_x) * num_pixels
logpZ1 = 0
factor_z = []
for (zi, mean, ln_var) in zs:
factor_z.append(zi.data.reshape(zi.shape[0], -1))
logpZ1 += cf.gaussian_nll(zi, mean, ln_var)
factor_z = xp.concatenate(factor_z, axis=1)
logpZ2 = cf.gaussian_nll(z, xp.zeros(z.shape), xp.zeros(z.shape)).data
# logpZ2 = cf.gaussian_nll(z, np.mean(z), np.log(np.var(z))).data
logpZ = (logpZ2 * 1 + logpZ1 * 1)
loss = 10 * b_norm + (logpZ - fw_ldt)
loss.backward()
optimizer.update()
training_step += 1
z_s.append(z.get())
b_s.append(cupy.asnumpy(b))
m_s.append(cupy.asnumpy(m.data))
loss_s.append(_float(loss))
logpZ_s.append(_float(logpZ))
logDet_s.append(_float(fw_ldt))
printr(
"Iteration {}: loss: {:.6f} - b_norm: {:.6f} - logpZ: {:.6f} - logpZ1: {:.6f} - logpZ2: {:.6f} - log_det: {:.6f} - logpX: {:.6f}\n"
.format(iteration + 1, _float(loss), _float(b_norm), _float(logpZ),
_float(logpZ1), _float(logpZ2), _float(fw_ldt),
_float(logpZ) - _float(fw_ldt)))
if iteration % 100 == 99:
print(cur_x.shape)
np.save(args.ckpt + '/' + str(j) + 'z.npy', z_s)
np.save(args.ckpt + '/' + str(j) + 'b.npy', b_s)
np.save(args.ckpt + '/' + str(j) + 'loss.npy', loss_s)
np.save(args.ckpt + '/' + str(j) + 'logpZ.npy', logpZ_s)
np.save(args.ckpt + '/' + str(j) + 'logDet.npy', logDet_s)
cur_x = make_uint8(cur_x[0].data, num_bins_x)
np.save(args.ckpt + '/' + str(j) + 'image.npy', cur_x)
x_PIL = Image.fromarray(cur_x)
x_PIL.save("./mask_imgs/trained.jpg")
np.save(args.ckpt + '/' + str(j) + 'm.npy', m_s)
# with encoder.reverse() as decoder:
# rx, _ = decoder.reverse_step(factor_z)
# rx_img = make_uint8(rx.data[0], num_bins_x)
# np.save(args.ckpt + '/'+str(j)+'res.npy', rx_img)
z_s = []
b_s = []
loss_s = []
logpZ_s = []
logDet_s = []
m_s = []
j += 1
epsilon.save(args.ckpt)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
image_size = (28, 28)
_, test = chainer.datasets.mnist.get_mnist()
images = []
labels = []
for entity in test:
image, label = entity
images.append(image)
labels.append(label)
labels = np.asarray(labels)
images = 255.0 * np.asarray(images).reshape((-1, ) + image_size + (1, ))
if hyperparams.num_image_channels != 1:
images = np.broadcast_to(images, (images.shape[0], ) + image_size +
(hyperparams.num_image_channels, ))
images = preprocess(images, hyperparams.num_bits_x)
# images = images[:200]
# labels = labels[:200]
sections = len(images) // 100
dataset_image = np.split(images, sections)
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
fig = plt.figure(figsize=(8, 8))
t_sne_inputs = []
with chainer.no_backprop_mode():
for n, image_batch in enumerate(dataset_image):
x = to_gpu(image_batch)
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
factorized_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
factorized_z.append(zi)
z = encoder.merge_factorized_z(factorized_z,
factor=hyperparams.squeeze_factor)
z = z.reshape((-1, hyperparams.num_image_channels * 28 * 28))
z = to_cpu(z)
t_sne_inputs.append(z)
t_sne_inputs = np.asanyarray(t_sne_inputs).reshape(
(-1, hyperparams.num_image_channels * 28 * 28))
print(t_sne_inputs.shape)
z_reduced = TSNE(n_components=2,
random_state=0).fit_transform(t_sne_inputs)
print(z_reduced.shape)
plt.scatter(z_reduced[:, 0],
z_reduced[:, 1],
c=labels,
s=1,
cmap="Spectral")
plt.colorbar()
plt.savefig("scatter.png")
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, hyperparams.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, hyperparams.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
dataset = glow.dataset.Dataset(images)
iterator = glow.dataset.Iterator(dataset, batch_size=2)
print(tabulate([["#image", len(dataset)]]))
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
total = args.num_steps + 2
fig = plt.figure(figsize=(4 * total, 4))
subplots = []
for n in range(total):
subplot = fig.add_subplot(1, total, n + 1)
subplots.append(subplot)
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
for data_indices in iterator:
x = to_gpu(dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
factorized_z_distribution, _ = encoder.forward_step(x)
factorized_z = []
for (zi, mean, ln_var) in factorized_z_distribution:
factorized_z.append(zi)
z = encoder.merge_factorized_z(factorized_z)
z_start = z[0]
z_end = z[1]
z_batch = [args.temperature * z_start]
for n in range(args.num_steps):
ratio = n / (args.num_steps - 1)
z_interp = ratio * z_end + (1.0 - ratio) * z_start
z_batch.append(args.temperature * z_interp)
z_batch.append(args.temperature * z_end)
for z, subplot in zip(z_batch, subplots):
z = z[None, ...]
rev_x, _ = decoder.reverse_step(z)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplot.imshow(rev_x_img, interpolation="none")
plt.pause(.01)
def main():
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
hyperparams = Hyperparameters(args.snapshot_path)
hyperparams.print()
num_bins_x = 2.0**hyperparams.num_bits_x
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
total = hyperparams.levels + 1
fig = plt.figure(figsize=(4 * total, 4))
subplots = []
for n in range(total):
subplot = fig.add_subplot(1, total, n + 1)
subplots.append(subplot)
def reverse_step(z, sampling=True):
if isinstance(z, list):
factorized_z = z
else:
factorized_z = encoder.factor_z(z)
assert len(factorized_z) == len(encoder.blocks)
out = None
sum_logdet = 0
for block, zi in zip(encoder.blocks[::-1], factorized_z[::-1]):
out, logdet = block.reverse_step(
out,
gaussian_eps=zi,
squeeze_factor=encoder.hyperparams.squeeze_factor,
sampling=sampling)
sum_logdet += logdet
return out, sum_logdet
with chainer.no_backprop_mode() and encoder.reverse() as decoder:
while True:
base_z = xp.random.normal(0,
args.temperature,
size=(
1,
3,
) + hyperparams.image_size,
dtype="float32")
factorized_z = encoder.factor_z(base_z)
rev_x, _ = decoder.reverse_step(factorized_z)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplots[0].imshow(rev_x_img, interpolation="none")
z = xp.copy(base_z)
factorized_z = encoder.factor_z(z)
for n in range(hyperparams.levels - 1):
factorized_z[n] = xp.random.normal(0,
args.temperature,
size=factorized_z[n].shape,
dtype="float32")
rev_x, _ = decoder.reverse_step(factorized_z)
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplots[1].imshow(rev_x_img, interpolation="none")
# for n in range(hyperparams.levels):
# z = xp.copy(base_z)
# factorized_z = encoder.factor_z(z)
# for m in range(n + 1):
# factorized_z[m] = xp.random.normal(
# 0,
# args.temperature,
# size=factorized_z[m].shape,
# dtype="float32")
# # factorized_z[m] = xp.zeros_like(factorized_z[m])
# out = None
# for k, (block, zi) in enumerate(
# zip(encoder.blocks[::-1], factorized_z[::-1])):
# sampling = False
# out, _ = block.reverse_step(
# out,
# gaussian_eps=zi,
# squeeze_factor=encoder.hyperparams.squeeze_factor,
# sampling=sampling)
# rev_x = out
# rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
# subplots[n + 1].imshow(rev_x_img, interpolation="none")
for n in range(hyperparams.levels):
z = xp.copy(base_z)
factorized_z = encoder.factor_z(z)
factorized_z[n] = xp.random.normal(0,
args.temperature,
size=factorized_z[n].shape,
dtype="float32")
factorized_z[n] = xp.zeros_like(factorized_z[n])
out = None
for k, (block, zi) in enumerate(
zip(encoder.blocks[::-1], factorized_z[::-1])):
sampling = False if k == hyperparams.levels - n - 1 else True
out, _ = block.reverse_step(
out,
gaussian_eps=zi,
squeeze_factor=encoder.hyperparams.squeeze_factor,
sampling=sampling)
rev_x = out
rev_x_img = make_uint8(rev_x.data[0], num_bins_x)
subplots[n + 1].imshow(rev_x_img, interpolation="none")
plt.pause(.01)
def main():
try:
os.mkdir(args.snapshot_path)
except:
pass
xp = np
using_gpu = args.gpu_device >= 0
if using_gpu:
cuda.get_device(args.gpu_device).use()
xp = cupy
num_bins_x = 2**args.num_bits_x
assert args.dataset_format in ["png", "npy"]
files = Path(args.dataset_path).glob("*.{}".format(args.dataset_format))
if args.dataset_format == "png":
images = []
for filepath in files:
image = np.array(Image.open(filepath)).astype("float32")
image = preprocess(image, args.num_bits_x)
images.append(image)
assert len(images) > 0
images = np.asanyarray(images)
elif args.dataset_format == "npy":
images = []
for filepath in files:
array = np.load(filepath).astype("float32")
array = preprocess(array, args.num_bits_x)
images.append(array)
assert len(images) > 0
num_files = len(images)
images = np.asanyarray(images)
images = images.reshape((num_files * images.shape[1], ) +
images.shape[2:])
else:
raise NotImplementedError
x_mean = np.mean(images)
x_var = np.var(images)
last_train_image = int((1 - args.validate_split) * len(images))
train_dataset = glow.dataset.Dataset(images[:last_train_image])
val_dataset = glow.dataset.Dataset(images[last_train_image:])
train_iterator = glow.dataset.Iterator(train_dataset,
batch_size=args.batch_size)
val_iterator = glow.dataset.Iterator(val_dataset,
batch_size=args.batch_size)
print(
tabulate([
["# train samples", len(train_dataset)],
["# validate samples", len(val_dataset)],
["overall mean", x_mean],
["overall var", x_var],
]))
hyperparams = Hyperparameters()
hyperparams.levels = args.levels
hyperparams.depth_per_level = args.depth_per_level
hyperparams.nn_hidden_channels = args.nn_hidden_channels
hyperparams.image_size = images.shape[2:]
hyperparams.num_bits_x = args.num_bits_x
hyperparams.lu_decomposition = args.lu_decomposition
hyperparams.squeeze_factor = args.squeeze_factor
hyperparams.save(args.snapshot_path)
hyperparams.print()
encoder = Glow(hyperparams, hdf5_path=args.snapshot_path)
if using_gpu:
encoder.to_gpu()
optimizer = Optimizer(encoder)
# Data dependent initialization
if encoder.need_initialize:
for batch_index, data_indices in enumerate(train_iterator):
x = to_gpu(train_dataset[data_indices])
encoder.initialize_actnorm_weights(x)
break
current_training_step = 0
num_pixels = 3 * hyperparams.image_size[0] * hyperparams.image_size[1]
# Training loop
for iteration in range(args.total_iteration):
# Training step
sum_loss = 0
sum_nll = 0
sum_kld = 0
start_time = time.time()
for batch_index, data_indices in enumerate(train_iterator):
x = to_gpu(train_dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
denom = math.log(2.0) * num_pixels
factorized_z_distribution, logdet = encoder.forward_step(x)
logdet -= math.log(num_bins_x) * num_pixels
kld = 0
negative_log_likelihood = 0
for (zi, mean, ln_var) in factorized_z_distribution:
negative_log_likelihood += cf.gaussian_nll(zi, mean, ln_var)
if args.regularize_z:
kld += cf.gaussian_kl_divergence(mean, ln_var)
loss = (negative_log_likelihood + kld) / args.batch_size - logdet
loss = loss / denom
encoder.cleargrads()
loss.backward()
optimizer.update(current_training_step)
current_training_step += 1
sum_loss += _float(loss)
sum_nll += _float(negative_log_likelihood) / args.batch_size
sum_kld += _float(kld) / args.batch_size
if (batch_index + 1) % 200 == 0:
print(
"Training Iteration {}: Batch {} / {} - loss: {:.8f} - nll: {:.8f} - kld: {:.8f} - log_det: {:.8f}"
.format(iteration + 1, batch_index + 1,
len(train_iterator), _float(loss),
_float(negative_log_likelihood) / args.batch_size,
_float(kld) / args.batch_size, _float(logdet)))
encoder.save(args.snapshot_path)
mean_log_likelihood = -sum_nll / len(train_iterator)
mean_kld = sum_kld / len(train_iterator)
elapsed_time = time.time() - start_time
print(
"\033[2KIteration {} Summary - training loss: {:.5f} - train log_likelihood: {:.5f} - train kld: {:.5f} - elapsed_time: {:.3f} min"
.format(iteration + 1, sum_loss / len(train_iterator),
mean_log_likelihood, mean_kld, elapsed_time / 60))
# Validation Step
sum_loss = 0
sum_nll = 0
sum_kld = 0
start_time = time.time()
with chainer.no_backprop_mode():
for batch_index, data_indices in enumerate(val_iterator):
x = to_gpu(val_dataset[data_indices])
x += xp.random.uniform(0, 1.0 / num_bins_x, size=x.shape)
denom = math.log(2.0) * num_pixels
factorized_z_distribution, logdet = encoder.forward_step(x)
logdet -= math.log(num_bins_x) * num_pixels
kld = 0
negative_log_likelihood = 0
for (zi, mean, ln_var) in factorized_z_distribution:
negative_log_likelihood += cf.gaussian_nll(
zi, mean, ln_var)
if args.regularize_z:
kld += cf.gaussian_kl_divergence(mean, ln_var)
loss = (negative_log_likelihood +
kld) / args.batch_size - logdet
loss = loss / denom
sum_loss += _float(loss)
sum_nll += _float(negative_log_likelihood) / args.batch_size
sum_kld += _float(kld) / args.batch_size
# Summary stats for iteration
mean_log_likelihood = -sum_nll / len(val_iterator)
mean_kld = sum_kld / len(val_iterator)
elapsed_time = time.time() - start_time
print(
"\033[2KIteration {} Summary - validation loss: {:.5f} - test log_likelihood: {:.5f} - test kld: {:.5f} - elapsed_time: {:.3f} min"
.format(iteration + 1, sum_loss / len(val_iterator),
mean_log_likelihood, mean_kld, elapsed_time / 60))
encoder.save(args.snapshot_path)
