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)