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 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 forward_and_reverse_output_shape(self,
in_channel,
data,
levels=3,
depth=4):
glow = Glow(in_channel, levels, depth)
z, logdet, eps = glow(data)
height, width = data.shape[2], data.shape[3]
"""
cifar example:
Level = 3
initial shape -> [4, 3, 32, 32]
iter 1 -> z: [4, 12, 16, 16] because of squeeze from outside the loop
iter 2 -> z: [4, 24, 8, 8] because of squeeze + split
iter 3 -> z: [4, 48, 4, 4] because of squeeze + split
"""
assert list(z.shape) == [4, in_channel * 4 * 2**(levels - 1), 4, 4]
assert list(logdet.shape) == [4] # because batch_size = 4
assert len(
eps
) == levels - 1 # because L = 3 and split is executed whenever < L, i.e 2 times in total
factor = 1
for e in eps:
factor *= 2
# example: first eps -> from iter 1 take z shape and divide channel by 2: [4, 12/2, 16, 16]
assert list(e.shape) == [
4, in_channel * factor, height / factor, width / factor
]
"""
In total depth * levels = 4 * 3 = 12, so we got 12 instances of actnorm, inconv and affinecoupling
Actnorm = 2 trainable parameters
Invconv = 3 trainable parameter
Affinecoupling = 6 trainable parameters (got 3 conv layers, each layer has weight + bias, so for all layers combined we get 6 in total)
Zeroconv = 4 (2 conv layers, each with weight + bias)
12 * (2+3+6) + 4= 136
"""
assert len(list(
glow.parameters())) == (levels * depth) * (2 + 3 + 6) + 4
for param in glow.parameters():
assert param.requires_grad
# reverse
# For cifar we expect z with level=3 to be of shape [4,48,4,4]
z = glow.reverse(z, eps)
assert list(z.shape) == [4, 3, 32, 32]
def test_generate_sample(self):
in_channel_cifar = 3
levels = 4
depth = 8
glow = Glow(in_channel_cifar, levels, depth)
x = torch.randn((2, 3, 32, 32))
x = generate(glow, 2, 'cpu', x.shape, levels)
# print(x)
print(x.shape)
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(args,kwargs):
output_folder = args.output_folder
model_name = args.model_name
with open(os.path.join(output_folder,'hparams.json')) as json_file:
hparams = json.load(json_file)
image_shape, num_classes, _, test_mnist = get_MNIST(False, hparams['dataroot'], hparams['download'])
test_loader = data.DataLoader(test_mnist, batch_size=32,
shuffle=False, num_workers=6,
drop_last=False)
x, y = test_loader.__iter__().__next__()
x = x.to(device)
model = Glow(image_shape, hparams['hidden_channels'], hparams['K'], hparams['L'], hparams['actnorm_scale'],
hparams['flow_permutation'], hparams['flow_coupling'], hparams['LU_decomposed'], num_classes,
hparams['learn_top'], hparams['y_condition'], False if 'logittransform' not in hparams else hparams['logittransform'],False if 'sn' not in hparams else hparams['sn'])
model.load_state_dict(torch.load(os.path.join(output_folder, model_name)))
model.set_actnorm_init()
model = model.to(device)
model = model.eval()
with torch.no_grad():
# ipdb.set_trace()
images = model(y_onehot=None, temperature=1, batch_size=32, reverse=True).cpu()
better_dup_images = model(y_onehot=None, temperature=1, z=model._last_z, reverse=True, use_last_split=True).cpu()
dup_images = model(y_onehot=None, temperature=1, z=model._last_z, reverse=True).cpu()
worse_dup_images = model(y_onehot=None, temperature=1, z=model._last_z, reverse=True).cpu()
l2_err = torch.pow((images - dup_images).view(images.shape[0], -1), 2).sum(-1).mean()
better_l2_err = torch.pow((images - better_dup_images).view(images.shape[0], -1), 2).sum(-1).mean()
worse_l2_err = torch.pow((images - worse_dup_images).view(images.shape[0], -1), 2).sum(-1).mean()
print(l2_err, better_l2_err, worse_l2_err)
plot_imgs([images, dup_images, better_dup_images, worse_dup_images], '_recons')
#
with torch.no_grad():
# ipdb.set_trace()
z, nll, y_logits = model(x, None)
better_dup_images = model(y_onehot=None, temperature=1, z=z, reverse=True, use_last_split=True).cpu()
plot_imgs([x, better_dup_images], '_data_recons2')
fpath = os.path.join(output_folder, '_recon_evoluation.png')
pad = run_recon_evolution(model, x, fpath)
def main(args):
# we're probably only be using 1 GPU, so this should be fine
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"running on {device}")
# set random seed for all
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
global best_loss
if(args.generate_samples):
print("generating samples")
# load data
# example for CIFAR-10 training:
train_set, test_set = get_dataloader(args.dataset, args.batch_size)
input_channels = channels_from_dataset(args.dataset)
print(f"amount of input channels: {input_channels}")
# instantiate model
# # baby network to make sure training script works
net = Glow(in_channels=input_channels,
depth=args.amt_flow_steps, levels=args.amt_levels, use_normalization=args.norm_method)
# code for rosalinty model
# net = RosGlow(input_channels, args.amt_flow_steps, args.amt_levels)
net = net.to(device)
print(f"training for {args.num_epochs} epochs.")
start_epoch = 0
# TODO: add functionality for loading checkpoints here
if args.resume:
print(f"resuming from checkpoint found in checkpoints/best_{args.dataset.lower()}.pth.tar.")
# raise error if no checkpoint directory is found
assert os.path.isdir("new_checkpoints")
checkpoint = torch.load(f"new_checkpoints/best_{args.dataset.lower()}.pth.tar")
net.load_state_dict(checkpoint["model"])
global best_loss
best_loss = checkpoint["test_loss"]
start_epoch = checkpoint["epoch"]
loss_function = FlowNLL().to(device)
optimizer = optim.Adam(net.parameters(), lr=float(args.lr))
# scheduler found in code, no mention in paper
# scheduler = sched.LambdaLR(
# optimizer, lambda s: min(1., s / args.warmup_iters))
# should we add a resume function here?
for epoch in range(start_epoch, start_epoch + args.num_epochs):
print(f"training epoch {epoch}")
train(net, train_set, device, optimizer, loss_function, epoch)
# how often do we want to test?
if (epoch % 10 == 0): # revert this to 10 once we know that this works
print(f"testing epoch {epoch}")
test(net, test_set, device, loss_function, epoch, args.generate_samples,
args.amt_levels, args.dataset, args.n_samples)
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 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)
f'Loss: {loss:.5f}; logP: {log_p:.5f}; logdet: {log_det:.5f}; lr: {warmup_lr:.7f}'
)
check_save(model_single, optimizer, args, z_sample, i)
except (KeyboardInterrupt, SystemExit):
check_save(model_single, optimizer, args, z_sample, i, save=True)
raise
if __name__ == '__main__':
args = parser.parse_args()
if len(args.load_path) > 0:
args.startiter = int(args.load_path[:-3].split('_')[-1])
print(args)
model_single = Glow(
1, args.n_flow, args.n_block, affine=args.affine, conv_lu=not args.no_lu
).cpu()
if len(args.load_path) > 0:
model_single.load_state_dict(torch.load(args.load_path, map_location=lambda storage, loc: storage))
model_single.initialize()
gc.collect()
torch.cuda.empty_cache()
model = model_single
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
if len(args.load_path) > 0:
optim_path = '/'.join(args.load_path.split('/')[:-1])
optimizer.load_state_dict(torch.load(os.path.join(optim_path, 'optimizer.pth'), map_location=lambda storage, loc: storage))
gc.collect()
model_path = args.model_path
checkpoint_path = os.path.join(model_path, args.checkpoint)
params_path = os.path.join(model_path, 'hparams.json')
with open(os.path.join(model_path, 'hparams.json')) as json_file:
hparams = json.load(json_file)
ds = check_dataset(args.dataset, args.dataroot, True, args.download)
ds2 = check_dataset(args.dataset2, args.dataroot, True, args.download)
image_shape, num_classes, train_dataset, test_dataset = ds
image_shape2, num_classes2, train_dataset_2, test_dataset_2 = ds2
model = Glow(image_shape, hparams['hidden_channels'], hparams['K'],
hparams['L'], hparams['actnorm_scale'],
hparams['flow_permutation'], hparams['flow_coupling'],
hparams['LU_decomposed'], num_classes, hparams['learn_top'],
hparams['y_condition'])
dic = torch.load(checkpoint_path)
if 'model' in dic.keys():
model.load_state_dict(dic["model"])
else:
model.load_state_dict(dic)
model.set_actnorm_init()
model = model.to(device)
model = model.eval()
if args.optim_type == "ADAM":
optim_default = partial(optim.Adam, lr=args.lr_test)
def main(args):
# torch.manual_seed(args.seed)
# Test loading and sampling
output_folder = os.path.join('results', args.name)
with open(os.path.join(output_folder, 'hparams.json')) as json_file:
hparams = json.load(json_file)
device = "cpu" if not torch.cuda.is_available() else "cuda:0"
image_shape = (hparams['patch_size'], hparams['patch_size'],
args.n_modalities)
num_classes = 1
print('Loading model...')
model = Glow(image_shape, hparams['hidden_channels'], hparams['K'],
hparams['L'], hparams['actnorm_scale'],
hparams['flow_permutation'], hparams['flow_coupling'],
hparams['LU_decomposed'], num_classes, hparams['learn_top'],
hparams['y_condition'])
model_chkpt = torch.load(
os.path.join(output_folder, 'checkpoints', args.model))
model.load_state_dict(model_chkpt['model'])
model.set_actnorm_init()
model = model.to(device)
# Build images
model.eval()
temperature = args.temperature
if args.steps is None: # automatically calculate step size if no step size
fig_dir = os.path.join(output_folder, 'stepnum_results')
if not os.path.exists(fig_dir):
os.mkdir(fig_dir)
print('No step size entered')
# Create sample of images to estimate chord length
with torch.no_grad():
mean, logs = model.prior(None, None)
z = gaussian_sample(mean, logs, temperature)
images_raw = model(z=z, temperature=temperature, reverse=True)
images_raw[torch.isnan(images_raw)] = 0.5
images_raw[torch.isinf(images_raw)] = 0.5
images_raw = torch.clamp(images_raw, -0.5, 0.5)
images_out = np.transpose(
np.squeeze(images_raw[:, args.step_modality, :, :].cpu().numpy()),
(1, 0, 2))
# Threshold images and compute covariances
if args.binary_data:
thresh = 0
else:
thresh = threshold_otsu(images_out)
images_bin = np.greater(images_out, thresh)
x_cov = two_point_correlation(images_bin, 0)
y_cov = two_point_correlation(images_bin, 1)
# Compute chord length
cov_avg = np.mean(np.mean(np.concatenate((x_cov, y_cov), axis=2),
axis=0),
axis=0)
N = 5
S20, _ = curve_fit(straight_line_at_origin(cov_avg[0]), range(0, N),
cov_avg[0:N])
l_pore = np.abs(cov_avg[0] / S20)
steps = int(l_pore)
print('Calculated step size: {}'.format(steps))
else:
print('Using user-entered step size {}...'.format(args.steps))
steps = args.steps
# Build desired number of volumes
for iter_vol in range(args.iter):
if args.iter == 1:
stack_dir = os.path.join(output_folder, 'image_stacks',
args.save_name)
print('Sampling images, saving to {}...'.format(args.save_name))
else:
stack_dir = os.path.join(
output_folder, 'image_stacks',
args.save_name + '_' + str(iter_vol).zfill(3))
print('Sampling images, saving to {}_'.format(args.save_name) +
str(iter_vol).zfill(3) + '...')
if not os.path.exists(stack_dir):
os.makedirs(stack_dir)
with torch.no_grad():
mean, logs = model.prior(None, None)
alpha = 1 - torch.reshape(torch.linspace(0, 1, steps=steps),
(-1, 1, 1, 1))
alpha = alpha.to(device)
num_imgs = int(np.ceil(hparams['patch_size'] / steps) + 1)
z = gaussian_sample(mean, logs, temperature)[:num_imgs, ...]
z = torch.cat([
alpha * z[i, ...] + (1 - alpha) * z[i + 1, ...]
for i in range(num_imgs - 1)
])
z = z[:hparams['patch_size'], ...]
images_raw = model(z=z, temperature=temperature, reverse=True)
images_raw[torch.isnan(images_raw)] = 0.5
images_raw[torch.isinf(images_raw)] = 0.5
images_raw = torch.clamp(images_raw, -0.5, 0.5)
# apply median filter to output
if args.med_filt is not None or args.binary_data:
for m in range(args.n_modalities):
if args.binary_data:
SE = ball(1)
else:
SE = ball(args.med_filt)
images_np = np.squeeze(images_raw[:, m, :, :].cpu().numpy())
images_filt = median_filter(images_np, footprint=SE)
# Erode binary images
if args.binary_data:
images_filt = np.greater(images_filt, 0)
SE = ball(1)
images_filt = 1.0 * binary_erosion(images_filt,
selem=SE) - 0.5
images_raw[:, m, :, :] = torch.tensor(images_filt,
device=device)
images1 = postprocess(images_raw).cpu()
images2 = postprocess(torch.transpose(images_raw, 0, 2)).cpu()
images3 = postprocess(torch.transpose(images_raw, 0, 3)).cpu()
# apply Otsu thresholding to output
if args.save_binary and not args.binary_data:
thresh = threshold_otsu(images1.numpy())
images1[images1 < thresh] = 0
images1[images1 > thresh] = 255
images2[images2 < thresh] = 0
images2[images2 > thresh] = 255
images3[images3 < thresh] = 0
images3[images3 > thresh] = 255
# # erode binary images by 1 px to correct for training image transformation
# if args.binary_data:
# images1 = np.greater(images1.numpy(), 127)
# images2 = np.greater(images2.numpy(), 127)
# images3 = np.greater(images3.numpy(), 127)
# images1 = 255*torch.tensor(1.0*np.expand_dims(binary_erosion(np.squeeze(images1), selem=np.ones((1,2,2))), 1))
# images2 = 255*torch.tensor(1.0*np.expand_dims(binary_erosion(np.squeeze(images2), selem=np.ones((2,1,2))), 1))
# images3 = 255*torch.tensor(1.0*np.expand_dims(binary_erosion(np.squeeze(images3), selem=np.ones((2,2,1))), 1))
# save video for each modality
for m in range(args.n_modalities):
if args.n_modalities > 1:
save_dir = os.path.join(stack_dir, 'modality{}'.format(m))
else:
save_dir = stack_dir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
write_video(images1[:, m, :, :], 'xy', hparams, save_dir)
write_video(images2[:, m, :, :], 'xz', hparams, save_dir)
write_video(images3[:, m, :, :], 'yz', hparams, save_dir)
print('Finished!')
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
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.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(dataset, dataroot, download, augment, batch_size, eval_batch_size,
epochs, saved_model, seed, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, learn_top,
y_condition, y_weight, max_grad_clip, max_grad_norm, lr, n_workers,
cuda, n_init_batches, warmup_steps, output_dir, saved_optimizer,
warmup, fresh, logittransform, gan, disc_lr, sn, flowgan, eval_every,
ld_on_samples, weight_gan, weight_prior, weight_logdet,
jac_reg_lambda, affine_eps, no_warm_up, optim_name, clamp, svd_every,
eval_only, no_actnorm, affine_scale_eps, actnorm_max_scale,
no_conv_actnorm, affine_max_scale, actnorm_eps, init_sample, no_split,
disc_arch, weight_entropy_reg, db):
check_manual_seed(seed)
ds = check_dataset(dataset, dataroot, augment, download)
image_shape, num_classes, train_dataset, test_dataset = ds
# Note: unsupported for now
multi_class = False
train_loader = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers,
drop_last=True)
test_loader = data.DataLoader(test_dataset,
batch_size=eval_batch_size,
shuffle=False,
num_workers=n_workers,
drop_last=False)
model = Glow(image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, num_classes,
learn_top, y_condition, logittransform, sn, affine_eps,
no_actnorm, affine_scale_eps, actnorm_max_scale,
no_conv_actnorm, affine_max_scale, actnorm_eps, no_split)
model = model.to(device)
if disc_arch == 'mine':
discriminator = mine.Discriminator(image_shape[-1])
elif disc_arch == 'biggan':
discriminator = cgan_models.Discriminator(
image_channels=image_shape[-1], conditional_D=False)
elif disc_arch == 'dcgan':
discriminator = DCGANDiscriminator(image_shape[0], 64, image_shape[-1])
elif disc_arch == 'inv':
discriminator = InvDiscriminator(
image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, num_classes,
learn_top, y_condition, logittransform, sn, affine_eps, no_actnorm,
affine_scale_eps, actnorm_max_scale, no_conv_actnorm,
affine_max_scale, actnorm_eps, no_split)
discriminator = discriminator.to(device)
D_optimizer = optim.Adam(filter(lambda p: p.requires_grad,
discriminator.parameters()),
lr=disc_lr,
betas=(.5, .99),
weight_decay=0)
if optim_name == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=lr,
betas=(.5, .99),
weight_decay=0)
elif optim_name == 'adamax':
optimizer = optim.Adamax(model.parameters(), lr=lr, weight_decay=5e-5)
if not no_warm_up:
lr_lambda = lambda epoch: min(1.0, (epoch + 1) / warmup)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lr_lambda)
iteration_fieldnames = [
'global_iteration', 'fid', 'sample_pad', 'train_bpd', 'eval_bpd',
'pad', 'batch_real_acc', 'batch_fake_acc', 'batch_acc'
]
iteration_logger = CSVLogger(fieldnames=iteration_fieldnames,
filename=os.path.join(output_dir,
'iteration_log.csv'))
iteration_fieldnames = [
'global_iteration', 'condition_num', 'max_sv', 'min_sv',
'inverse_condition_num', 'inverse_max_sv', 'inverse_min_sv'
]
svd_logger = CSVLogger(fieldnames=iteration_fieldnames,
filename=os.path.join(output_dir, 'svd_log.csv'))
#
test_iter = test_loader.__iter__()
N_inception = 1000
x_real_inception = torch.cat([
test_iter.__next__()[0].to(device)
for _ in range(N_inception // args.batch_size + 1)
], 0)[:N_inception]
x_real_inception = x_real_inception + .5
x_for_recon = test_iter.__next__()[0].to(device)
def gan_step(engine, batch):
assert not y_condition
if 'iter_ind' in dir(engine):
engine.iter_ind += 1
else:
engine.iter_ind = -1
losses = {}
model.train()
discriminator.train()
x, y = batch
x = x.to(device)
def run_noised_disc(discriminator, x):
x = uniform_binning_correction(x)[0]
return discriminator(x)
real_acc = fake_acc = acc = 0
if weight_gan > 0:
fake = generate_from_noise(model, x.size(0), clamp=clamp)
D_real_scores = run_noised_disc(discriminator, x.detach())
D_fake_scores = run_noised_disc(discriminator, fake.detach())
ones_target = torch.ones((x.size(0), 1), device=x.device)
zeros_target = torch.zeros((x.size(0), 1), device=x.device)
D_real_accuracy = torch.sum(
torch.round(F.sigmoid(D_real_scores)) ==
ones_target).float() / ones_target.size(0)
D_fake_accuracy = torch.sum(
torch.round(F.sigmoid(D_fake_scores)) ==
zeros_target).float() / zeros_target.size(0)
D_real_loss = F.binary_cross_entropy_with_logits(
D_real_scores, ones_target)
D_fake_loss = F.binary_cross_entropy_with_logits(
D_fake_scores, zeros_target)
D_loss = (D_real_loss + D_fake_loss) / 2
gp = gradient_penalty(
x.detach(), fake.detach(),
lambda _x: run_noised_disc(discriminator, _x))
D_loss_plus_gp = D_loss + 10 * gp
D_optimizer.zero_grad()
D_loss_plus_gp.backward()
D_optimizer.step()
# Train generator
fake = generate_from_noise(model,
x.size(0),
clamp=clamp,
guard_nans=False)
G_loss = F.binary_cross_entropy_with_logits(
run_noised_disc(discriminator, fake),
torch.ones((x.size(0), 1), device=x.device))
# Trace
real_acc = D_real_accuracy.item()
fake_acc = D_fake_accuracy.item()
acc = .5 * (D_fake_accuracy.item() + D_real_accuracy.item())
z, nll, y_logits, (prior, logdet) = model.forward(x,
None,
return_details=True)
train_bpd = nll.mean().item()
loss = 0
if weight_gan > 0:
loss = loss + weight_gan * G_loss
if weight_prior > 0:
loss = loss + weight_prior * -prior.mean()
if weight_logdet > 0:
loss = loss + weight_logdet * -logdet.mean()
if weight_entropy_reg > 0:
_, _, _, (sample_prior,
sample_logdet) = model.forward(fake,
None,
return_details=True)
# notice this is actually "decreasing" sample likelihood.
loss = loss + weight_entropy_reg * (sample_prior.mean() +
sample_logdet.mean())
# Jac Reg
if jac_reg_lambda > 0:
# Sample
x_samples = generate_from_noise(model,
args.batch_size,
clamp=clamp).detach()
x_samples.requires_grad_()
z = model.forward(x_samples, None, return_details=True)[0]
other_zs = torch.cat([
split._last_z2.view(x.size(0), -1)
for split in model.flow.splits
], -1)
all_z = torch.cat([other_zs, z.view(x.size(0), -1)], -1)
sample_foward_jac = compute_jacobian_regularizer(x_samples,
all_z,
n_proj=1)
_, c2, h, w = model.prior_h.shape
c = c2 // 2
zshape = (batch_size, c, h, w)
randz = torch.randn(zshape).to(device)
randz = torch.autograd.Variable(randz, requires_grad=True)
images = model(z=randz,
y_onehot=None,
temperature=1,
reverse=True,
batch_size=0)
other_zs = [split._last_z2 for split in model.flow.splits]
all_z = [randz] + other_zs
sample_inverse_jac = compute_jacobian_regularizer_manyinputs(
all_z, images, n_proj=1)
# Data
x.requires_grad_()
z = model.forward(x, None, return_details=True)[0]
other_zs = torch.cat([
split._last_z2.view(x.size(0), -1)
for split in model.flow.splits
], -1)
all_z = torch.cat([other_zs, z.view(x.size(0), -1)], -1)
data_foward_jac = compute_jacobian_regularizer(x, all_z, n_proj=1)
_, c2, h, w = model.prior_h.shape
c = c2 // 2
zshape = (batch_size, c, h, w)
z.requires_grad_()
images = model(z=z,
y_onehot=None,
temperature=1,
reverse=True,
batch_size=0)
other_zs = [split._last_z2 for split in model.flow.splits]
all_z = [z] + other_zs
data_inverse_jac = compute_jacobian_regularizer_manyinputs(
all_z, images, n_proj=1)
# loss = loss + jac_reg_lambda * (sample_foward_jac + sample_inverse_jac )
loss = loss + jac_reg_lambda * (sample_foward_jac +
sample_inverse_jac +
data_foward_jac + data_inverse_jac)
if not eval_only:
optimizer.zero_grad()
loss.backward()
if not db:
assert max_grad_clip == max_grad_norm == 0
if max_grad_clip > 0:
torch.nn.utils.clip_grad_value_(model.parameters(),
max_grad_clip)
if max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(),
max_grad_norm)
# Replace NaN gradient with 0
for p in model.parameters():
if p.requires_grad and p.grad is not None:
g = p.grad.data
g[g != g] = 0
optimizer.step()
if engine.iter_ind % 100 == 0:
with torch.no_grad():
fake = generate_from_noise(model, x.size(0), clamp=clamp)
z = model.forward(fake, None, return_details=True)[0]
print("Z max min")
print(z.max().item(), z.min().item())
if (fake != fake).float().sum() > 0:
title = 'NaNs'
else:
title = "Good"
grid = make_grid((postprocess(fake.detach().cpu(), dataset)[:30]),
nrow=6).permute(1, 2, 0)
plt.figure(figsize=(10, 10))
plt.imshow(grid)
plt.axis('off')
plt.title(title)
plt.savefig(
os.path.join(output_dir, f'sample_{engine.iter_ind}.png'))
if engine.iter_ind % eval_every == 0:
def check_all_zero_except_leading(x):
return x % 10**np.floor(np.log10(x)) == 0
if engine.iter_ind == 0 or check_all_zero_except_leading(
engine.iter_ind):
torch.save(
model.state_dict(),
os.path.join(output_dir, f'ckpt_sd_{engine.iter_ind}.pt'))
model.eval()
with torch.no_grad():
# Plot recon
fpath = os.path.join(output_dir, '_recon',
f'recon_{engine.iter_ind}.png')
sample_pad = run_recon_evolution(
model,
generate_from_noise(model, args.batch_size,
clamp=clamp).detach(), fpath)
print(
f"Iter: {engine.iter_ind}, Recon Sample PAD: {sample_pad}")
pad = run_recon_evolution(model, x_for_recon, fpath)
print(f"Iter: {engine.iter_ind}, Recon PAD: {pad}")
pad = pad.item()
sample_pad = sample_pad.item()
# Inception score
sample = torch.cat([
generate_from_noise(model, args.batch_size, clamp=clamp)
for _ in range(N_inception // args.batch_size + 1)
], 0)[:N_inception]
sample = sample + .5
if (sample != sample).float().sum() > 0:
print("Sample NaNs")
raise
else:
fid = run_fid(x_real_inception.clamp_(0, 1),
sample.clamp_(0, 1))
print(f'fid: {fid}, global_iter: {engine.iter_ind}')
# Eval BPD
eval_bpd = np.mean([
model.forward(x.to(device), None,
return_details=True)[1].mean().item()
for x, _ in test_loader
])
stats_dict = {
'global_iteration': engine.iter_ind,
'fid': fid,
'train_bpd': train_bpd,
'pad': pad,
'eval_bpd': eval_bpd,
'sample_pad': sample_pad,
'batch_real_acc': real_acc,
'batch_fake_acc': fake_acc,
'batch_acc': acc
}
iteration_logger.writerow(stats_dict)
plot_csv(iteration_logger.filename)
model.train()
if engine.iter_ind + 2 % svd_every == 0:
model.eval()
svd_dict = {}
ret = utils.computeSVDjacobian(x_for_recon, model)
D_for, D_inv = ret['D_for'], ret['D_inv']
cn = float(D_for.max() / D_for.min())
cn_inv = float(D_inv.max() / D_inv.min())
svd_dict['global_iteration'] = engine.iter_ind
svd_dict['condition_num'] = cn
svd_dict['max_sv'] = float(D_for.max())
svd_dict['min_sv'] = float(D_for.min())
svd_dict['inverse_condition_num'] = cn_inv
svd_dict['inverse_max_sv'] = float(D_inv.max())
svd_dict['inverse_min_sv'] = float(D_inv.min())
svd_logger.writerow(svd_dict)
# plot_utils.plot_stability_stats(output_dir)
# plot_utils.plot_individual_figures(output_dir, 'svd_log.csv')
model.train()
if eval_only:
sys.exit()
# Dummy
losses['total_loss'] = torch.mean(nll).item()
return losses
def eval_step(engine, batch):
model.eval()
x, y = batch
x = x.to(device)
with torch.no_grad():
if y_condition:
y = y.to(device)
z, nll, y_logits = model(x, y)
losses = compute_loss_y(nll,
y_logits,
y_weight,
y,
multi_class,
reduction='none')
else:
z, nll, y_logits = model(x, None)
losses = compute_loss(nll, reduction='none')
return losses
trainer = Engine(gan_step)
# else:
# trainer = Engine(step)
checkpoint_handler = ModelCheckpoint(output_dir,
'glow',
save_interval=5,
n_saved=1,
require_empty=False)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {
'model': model,
'optimizer': optimizer
})
monitoring_metrics = ['total_loss']
RunningAverage(output_transform=lambda x: x['total_loss']).attach(
trainer, 'total_loss')
evaluator = Engine(eval_step)
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(lambda x, y: torch.mean(x),
output_transform=lambda x:
(x['total_loss'], torch.empty(x['total_loss'].shape[0]))).attach(
evaluator, 'total_loss')
if y_condition:
monitoring_metrics.extend(['nll'])
RunningAverage(output_transform=lambda x: x['nll']).attach(
trainer, 'nll')
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(lambda x, y: torch.mean(x),
output_transform=lambda x:
(x['nll'], torch.empty(x['nll'].shape[0]))).attach(
evaluator, 'nll')
pbar = ProgressBar()
pbar.attach(trainer, metric_names=monitoring_metrics)
# load pre-trained model if given
if saved_model:
print("Loading...")
print(saved_model)
loaded = torch.load(saved_model)
# if 'Glow' in str(type(loaded)):
# model = loaded
# else:
# raise
# # if 'Glow' in str(type(loaded)):
# # loaded = loaded.state_dict()
model.load_state_dict(loaded)
model.set_actnorm_init()
if saved_optimizer:
optimizer.load_state_dict(torch.load(saved_optimizer))
file_name, ext = os.path.splitext(saved_model)
resume_epoch = int(file_name.split('_')[-1])
@trainer.on(Events.STARTED)
def resume_training(engine):
engine.state.epoch = resume_epoch
engine.state.iteration = resume_epoch * len(
engine.state.dataloader)
@trainer.on(Events.STARTED)
def init(engine):
if saved_model:
return
model.train()
print("Initializing Actnorm...")
init_batches = []
init_targets = []
if n_init_batches == 0:
model.set_actnorm_init()
return
with torch.no_grad():
if init_sample:
generate_from_noise(model,
args.batch_size * args.n_init_batches)
else:
for batch, target in islice(train_loader, None,
n_init_batches):
init_batches.append(batch)
init_targets.append(target)
init_batches = torch.cat(init_batches).to(device)
assert init_batches.shape[0] == n_init_batches * batch_size
if y_condition:
init_targets = torch.cat(init_targets).to(device)
else:
init_targets = None
model(init_batches, init_targets)
@trainer.on(Events.EPOCH_COMPLETED)
def evaluate(engine):
evaluator.run(test_loader)
if not no_warm_up:
scheduler.step()
metrics = evaluator.state.metrics
losses = ', '.join(
[f"{key}: {value:.2f}" for key, value in metrics.items()])
print(f'Validation Results - Epoch: {engine.state.epoch} {losses}')
timer = Timer(average=True)
timer.attach(trainer,
start=Events.EPOCH_STARTED,
resume=Events.ITERATION_STARTED,
pause=Events.ITERATION_COMPLETED,
step=Events.ITERATION_COMPLETED)
@trainer.on(Events.EPOCH_COMPLETED)
def print_times(engine):
pbar.log_message(
f'Epoch {engine.state.epoch} done. Time per batch: {timer.value():.3f}[s]'
)
timer.reset()
trainer.run(train_loader, epochs)
)
parser.add_argument("--affine",
action="store_true",
help="use affine coupling instead of additive")
parser.add_argument("--n_flow",
default=32,
type=int,
help="number of flows in each block")
parser.add_argument("--n_bits", default=5, type=int, help="number of bits")
parser.add_argument("--img_size", default=64, type=int)
parser.add_argument("--batch_size", default=18, type=int)
args = parser.parse_args()
model_single = Glow(3,
args.n_flow,
args.n_block,
affine=args.affine,
conv_lu=not args.no_lu)
model = torch.nn.DataParallel(model_single)
model = model.to(device)
ckt = torch.load('savemodel/checkpoint_18.tar')
model.load_state_dict(ckt['net'])
dataset = CelebALoader('./', args.img_size)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True)
with_attr_cnt = 0
wo_attr_cnt = 0
with_attr_z = None
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(
dataset,
dataset2,
dataroot,
download,
augment,
batch_size,
eval_batch_size,
nlls_batch_size,
epochs,
nb_step,
saved_model,
seed,
hidden_channels,
K,
L,
actnorm_scale,
flow_permutation,
flow_coupling,
LU_decomposed,
learn_top,
y_condition,
y_weight,
max_grad_clip,
max_grad_norm,
lr,
lr_test,
n_workers,
cuda,
n_init_batches,
output_dir,
saved_optimizer,
warmup,
every_epoch,
):
device = "cpu" if (not torch.cuda.is_available() or not cuda) else "cuda:0"
check_manual_seed(seed)
ds = check_dataset(dataset, dataroot, augment, download)
ds2 = check_dataset(dataset2, dataroot, augment, download)
image_shape, num_classes, train_dataset, test_dataset = ds
image_shape2, num_classes2, train_dataset_2, test_dataset_2 = ds2
assert(image_shape == image_shape2)
data1 = []
data2 = []
for k in range(nlls_batch_size):
dataaux, targetaux = test_dataset[k]
data1.append(dataaux)
dataaux, targetaux = test_dataset_2[k]
data2.append(dataaux)
# Note: unsupported for now
multi_class = False
train_loader = data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers,
drop_last=True,
)
test_loader = data.DataLoader(
test_dataset,
batch_size=eval_batch_size,
shuffle=False,
num_workers=n_workers,
drop_last=False,
)
model = Glow(
image_shape,
hidden_channels,
K,
L,
actnorm_scale,
flow_permutation,
flow_coupling,
LU_decomposed,
num_classes,
learn_top,
y_condition,
)
model = model.to(device)
optimizer = optim.Adamax(model.parameters(), lr=lr, weight_decay=5e-5)
lr_lambda = lambda epoch: min(1.0, (epoch + 1) / warmup) # noqa
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
def step(engine, batch):
model.train()
optimizer.zero_grad()
x, y = batch
x = x.to(device)
if y_condition:
y = y.to(device)
z, nll, y_logits = model(x, y)
losses = compute_loss_y(nll, y_logits, y_weight, y, multi_class)
else:
z, nll, y_logits = model(x, None)
losses = compute_loss(nll)
losses["total_loss"].backward()
if max_grad_clip > 0:
torch.nn.utils.clip_grad_value_(model.parameters(), max_grad_clip)
if max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
return losses
def eval_step(engine, batch):
model.eval()
x, y = batch
x = x.to(device)
with torch.no_grad():
if y_condition:
y = y.to(device)
z, nll, y_logits = model(x, y)
losses = compute_loss_y(
nll, y_logits, y_weight, y, multi_class, reduction="none"
)
else:
z, nll, y_logits = model(x, None)
losses = compute_loss(nll, reduction="none")
return losses
trainer = Engine(step)
checkpoint_handler = ModelCheckpoint(
output_dir, "glow", n_saved=2, require_empty=False
)
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
checkpoint_handler,
{"model": model, "optimizer": optimizer},
)
monitoring_metrics = ["total_loss"]
RunningAverage(output_transform=lambda x: x["total_loss"]).attach(
trainer, "total_loss"
)
evaluator = Engine(eval_step)
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(
lambda x, y: torch.mean(x),
output_transform=lambda x: (
x["total_loss"],
torch.empty(x["total_loss"].shape[0]),
),
).attach(evaluator, "total_loss")
if y_condition:
monitoring_metrics.extend(["nll"])
RunningAverage(output_transform=lambda x: x["nll"]).attach(trainer, "nll")
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(
lambda x, y: torch.mean(x),
output_transform=lambda x: (x["nll"], torch.empty(x["nll"].shape[0])),
).attach(evaluator, "nll")
pbar = ProgressBar()
pbar.attach(trainer, metric_names=monitoring_metrics)
# load pre-trained model if given
if saved_model:
model.load_state_dict(torch.load(saved_model)['model'])
model.set_actnorm_init()
if saved_optimizer:
optimizer.load_state_dict(torch.load(saved_optimizer)['opt'])
file_name, ext = os.path.splitext(saved_model)
resume_epoch = int(file_name.split("_")[-1])/1e3
@trainer.on(Events.STARTED)
def resume_training(engine):
engine.state.epoch = resume_epoch
engine.state.iteration = resume_epoch * len(engine.state.dataloader)
@trainer.on(Events.STARTED)
def init(engine):
model.train()
init_batches = []
init_targets = []
with torch.no_grad():
print(train_loader)
for batch, target in islice(train_loader, None, n_init_batches):
init_batches.append(batch)
init_targets.append(target)
init_batches = torch.cat(init_batches).to(device)
assert init_batches.shape[0] == n_init_batches * batch_size
if y_condition:
init_targets = torch.cat(init_targets).to(device)
else:
init_targets = None
model(init_batches, init_targets)
@trainer.on(Events.EPOCH_COMPLETED)
def evaluate(engine):
evaluator.run(test_loader)
scheduler.step()
metrics = evaluator.state.metrics
losses = ", ".join([f"{key}: {value:.2f}" for key, value in metrics.items()])
print(f"Validation Results - Epoch: {engine.state.epoch} {losses}")
timer = Timer(average=True)
timer.attach(
trainer,
start=Events.EPOCH_STARTED,
resume=Events.ITERATION_STARTED,
pause=Events.ITERATION_COMPLETED,
step=Events.ITERATION_COMPLETED,
)
@trainer.on(Events.EPOCH_COMPLETED)
def print_times(engine):
pbar.log_message(
f"Epoch {engine.state.epoch} done. Time per batch: {timer.value():.3f}[s]"
)
timer.reset()
# @trainer.on(Events.EPOCH_COMPLETED)
# def eval_likelihood(engine):
# global_nlls(output_dir, engine.state.epoch, data1, data2, model, dataset1_name = dataset, dataset2_name = dataset2, nb_step = nb_step, every_epoch = every_epoch, optim_default = partial(optim.SGD, lr=1e-5, momentum = 0.))
trainer.run(train_loader, epochs)
def main(args):
seed_list = [int(item) for item in args.seed.split(',')]
for seed in seed_list:
device = torch.device("cuda")
experiment_folder = args.experiment_folder + '/' + str(seed) + '/'
print(experiment_folder)
#model_name = 'glow_checkpoint_'+ str(args.chk)+'.pth'
for thing in os.listdir(experiment_folder):
if 'best' in thing:
model_name = thing
print(model_name)
random.seed(seed)
torch.manual_seed(seed)
np.random.seed(seed)
torch.random.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
with open(experiment_folder + 'hparams.json') as json_file:
hparams = json.load(json_file)
image_shape = (32, 32, 3)
if hparams['y_condition']:
num_classes = 2
num_domains = 0
elif hparams['d_condition']:
num_classes = 10
num_domains = 0
elif hparams['yd_condition']:
num_classes = 2
num_domains = 10
else:
num_classes = 2
num_domains = 0
model = Glow(image_shape, hparams['hidden_channels'], hparams['K'],
hparams['L'], hparams['actnorm_scale'],
hparams['flow_permutation'], hparams['flow_coupling'],
hparams['LU_decomposed'], num_classes, num_domains,
hparams['learn_top'], hparams['y_condition'],
hparams['extra_condition'], hparams['sp_condition'],
hparams['d_condition'], hparams['yd_condition'])
print('loading model')
model.load_state_dict(torch.load(experiment_folder + model_name))
model.set_actnorm_init()
model = model.to(device)
model = model.eval()
if hparams['y_condition']:
print('y_condition')
def sample(model, temp=args.temperature):
with torch.no_grad():
if hparams['y_condition']:
print("extra", hparams['extra_condition'])
y = torch.eye(num_classes)
y = torch.cat(1000 * [y])
print(y.size())
y_0 = y[::2, :].to(
device) # number hardcoded in model for now
y_1 = y[1::2, :].to(device)
print(y_0.size())
print(y_0)
print(y_1)
print(y_1.size())
images0 = model(z=None,
y_onehot=y_0,
temperature=temp,
reverse=True,
batch_size=1000)
images1 = model(z=None,
y_onehot=y_1,
temperature=temp,
reverse=True,
batch_size=1000)
return images0, images1
images0, images1 = sample(model)
os.makedirs(experiment_folder + 'generations/Uninfected',
exist_ok=True)
os.makedirs(experiment_folder + 'generations/Parasitized',
exist_ok=True)
for i in range(images0.size(0)):
torchvision.utils.save_image(
images0[i, :, :, :], experiment_folder +
'generations/Uninfected/sample_{}.png'.format(i))
torchvision.utils.save_image(
images1[i, :, :, :], experiment_folder +
'generations/Parasitized/sample_{}.png'.format(i))
images_concat0 = torchvision.utils.make_grid(images0[:64, :, :, :],
nrow=int(64**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(images_concat0,
experiment_folder + '/uninfected.png')
images_concat1 = torchvision.utils.make_grid(images1[:64, :, :, :],
nrow=int(64**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(
images_concat1, experiment_folder + '/parasitized.png')
elif hparams['d_condition']:
print('d_cond')
def sample_d(model, idx, batch_size=1000, temp=args.temperature):
with torch.no_grad():
if hparams['d_condition']:
y_0 = torch.zeros([batch_size, 10], device='cuda:0')
y_0[:, idx] = torch.ones(batch_size)
y_0.to(device)
print(y_0)
# y_1 = torch.zeros([batch_size, 201], device='cuda:0')
# y_1[:, 157] = torch.ones(batch_size)
# y_1.to(device)
# y = torch.eye(num_classes)
# y = torch.cat(1000 * [y])
# print(y.size())
# y_0 = y[::2, :].to(device) # number hardcoded in model for now
# y_1 = y[1::2, :].to(device)
# print(y_0.size())
# print(y_0)
# print(y_1)
# print(y_1.size())
images0 = model(z=None,
y_onehot=y_0,
temperature=temp,
reverse=True,
batch_size=1000)
# images1 = model(z=None, y_onehot=y_1, temperature=1.0, reverse=True, batch_size=1000)
return images0
for idx, dom in enumerate(["C116P77ThinF", "C132P93ThinF", "C137P98ThinF", "C180P141NThinF", "C182P143NThinF", \
"C184P145ThinF", "C39P4thinF", 'C59P20thinF', "C68P29N", "C99P60ThinF"]):
images0 = sample_d(model, idx)
os.makedirs(experiment_folder + 'generations/' + dom +
'/Uninfected/',
exist_ok=True)
os.makedirs(experiment_folder + 'generations/' + dom +
'/Parasitized/',
exist_ok=True)
# os.makedirs(experiment_folder + 'generations/C59P20thinF/Uninfected/', exist_ok=True)
# os.makedirs(experiment_folder + 'generations/C59P20thinF/Parasitized/', exist_ok=True)
for i in range(images0.size(0)):
torchvision.utils.save_image(
images0[i, :, :, :],
experiment_folder + 'generations/' + dom +
'/Uninfected/sample_{}.png'.format(i))
#torchvision.utils.save_image(images1[i, :, :, :], experiment_folder + 'generations/C59P20thinF/Parasitized/sample_{}.png'.format(i))
images_concat0 = torchvision.utils.make_grid(
images0[:25, :, :, :],
nrow=int(25**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(images_concat0,
experiment_folder + dom + '.png')
# images_concat1 = torchvision.utils.make_grid(images1[:64,:,:,:], nrow=int(64 ** 0.5), padding=2, pad_value=255)
# torchvision.utils.save_image(images_concat1, experiment_folder + 'C59P20thinF.png')
elif hparams['yd_condition']:
def sample_YD(model, idx, batch_size=1000, temp=args.temperature):
with torch.no_grad():
if hparams['yd_condition']:
y_0 = torch.zeros([batch_size, 12], device='cuda:0')
y_0[:, 0] = torch.ones(batch_size)
y_0[:, idx + 2] = torch.ones(batch_size)
y_0.to(device)
print(y_0)
y_1 = torch.zeros([batch_size, 12], device='cuda:0')
y_1[:, 1] = torch.ones(batch_size)
y_1[:, idx + 2] = torch.ones(batch_size)
y_1.to(device)
print(y_1)
images0 = model(z=None,
y_onehot=y_0,
temperature=temp,
reverse=True,
batch_size=1000)
images1 = model(z=None,
y_onehot=y_1,
temperature=temp,
reverse=True,
batch_size=1000)
return images0, images1
def sample_DD(model, idx, batch_size=1000, temp=args.temperature):
with torch.no_grad():
if hparams['yd_condition']:
y_1 = torch.zeros([batch_size, 20], device='cuda:0')
y_1[:, idx] = torch.ones(batch_size)
y_1.to(device)
print(y_1)
y_0 = torch.zeros([batch_size, 20], device='cuda:0')
y_0[:, idx + 10] = torch.ones(batch_size)
y_0.to(device)
print(y_0)
images0 = model(z=None,
y_onehot=y_0,
temperature=temp,
reverse=True,
batch_size=1000)
images1 = model(z=None,
y_onehot=y_1,
temperature=temp,
reverse=True,
batch_size=1000)
return images0, images1
for idx, dom in enumerate(
["C116P77ThinF", "C132P93ThinF", "C137P98ThinF", "C180P141NThinF", "C182P143NThinF", \
"C184P145ThinF", "C39P4thinF", 'C59P20thinF', "C68P29N", "C99P60ThinF"]):
images0, images1 = sample_YD(model, idx)
os.makedirs(experiment_folder + 'generations/' + dom +
'/Uninfected/',
exist_ok=True)
os.makedirs(experiment_folder + 'generations/' + dom +
'/Parasitized/',
exist_ok=True)
for i in range(images0.size(0)):
torchvision.utils.save_image(
images0[i, :, :, :],
experiment_folder + 'generations/' + dom +
'/Uninfected/sample_{}.png'.format(i))
torchvision.utils.save_image(
images1[i, :, :, :],
experiment_folder + 'generations/' + dom +
'/Parasitized/sample_{}.png'.format(i))
images_concat0 = torchvision.utils.make_grid(
images0[:64, :, :, :],
nrow=int(64**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(
images_concat0, experiment_folder + dom +
str(args.temperature) + '_uninfected.png')
images_concat1 = torchvision.utils.make_grid(
images1[:64, :, :, :],
nrow=int(64**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(
images_concat1, experiment_folder + dom +
str(args.temperature) + '_parasitized.png')
else:
def sample(model, temp=args.temperature):
with torch.no_grad():
images = model(z=None,
y_onehot=None,
temperature=temp,
reverse=True,
batch_size=1000)
return images
images = sample(model)
os.makedirs('unconditioned/' + str(seed) + '/generations/' +
experiment_folder[:-3],
exist_ok=True)
for i in range(images.size(0)):
torchvision.utils.save_image(
images[i, :, :, :],
'unconditioned/' + str(seed) + '/generations/' +
experiment_folder[:-2] + 'sample_{}.png'.format(i))
images_concat = torchvision.utils.make_grid(images[:64, :, :, :],
nrow=int(64**0.5),
padding=2,
pad_value=255)
torchvision.utils.save_image(
images_concat, 'unconditioned/' + str(seed) + '/' +
experiment_folder[:-3] + '.png')
# ipdb.set_trace()
with torch.no_grad():
iss, _, _, acts_real = inception_score(x_is, cuda=True, batch_size=32, resize=True, splits=10, return_preds=True)
print(iss)
# Model samples
coupling = 'affine'
sn_ = 0
loss_ = 'mle'
# for (coupling, sn_) in product(['affine','additive'], [0,1]):
exp_dir = f'/scratch/gobi2/wangkuan/glow/rebuttal-guess/{loss_}-128-32-{coupling}-{sn_}'
with open(os.path.join(exp_dir, 'hparams.json'), 'r') as f:
params = json.load(f)
locals().update(params)
model = Glow(image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, num_classes,
learn_top, y_condition,logittransform,sn)
model = model.to(device)
# ipdb.set_trace()
def generate_from_noise(batch_size):
_, c2, h, w = model.prior_h.shape
c = c2 // 2
zshape = (batch_size, c, h, w)
randz = torch.randn(zshape).to(device)
randz = torch.autograd.Variable(randz, requires_grad=True)
images = model(z= randz, y_onehot=None, temperature=1, reverse=True,batch_size=batch_size)
return images
# ipdb.set_trace()
iteration_fieldnames = ['global_iteration', 'fid']
def main(dataset, dataroot, download, augment, batch_size, eval_batch_size,
epochs, saved_model, seed, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, learn_top,
y_condition, y_weight, max_grad_clip, max_grad_norm, lr, n_workers,
cuda, n_init_batches, warmup_steps, output_dir, saved_optimizer,
fresh):
device = 'cpu' if (not torch.cuda.is_available() or not cuda) else 'cuda:0'
check_manual_seed(seed)
ds = check_dataset(dataset, dataroot, augment, download)
image_shape, num_classes, train_dataset, test_dataset = ds
# Note: unsupported for now
multi_class = False
train_loader = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers,
drop_last=True)
test_loader = data.DataLoader(test_dataset,
batch_size=eval_batch_size,
shuffle=False,
num_workers=n_workers,
drop_last=False)
model = Glow(image_shape, hidden_channels, K, L, actnorm_scale,
flow_permutation, flow_coupling, LU_decomposed, num_classes,
learn_top, y_condition)
model = model.to(device)
optimizer = optim.Adamax(model.parameters(), lr=lr, weight_decay=5e-5)
def step(engine, batch):
model.train()
optimizer.zero_grad()
x, y = batch
x = x.to(device)
if y_condition:
y = y.to(device)
z, nll, y_logits = model(x, y)
losses = compute_loss_y(nll, y_logits, y_weight, y, multi_class)
else:
z, nll, y_logits = model(x, None)
losses = compute_loss(nll)
losses['total_loss'].backward()
if max_grad_clip > 0:
torch.nn.utils.clip_grad_value_(model.parameters(), max_grad_clip)
if max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
return losses
def eval_step(engine, batch):
model.eval()
x, y = batch
x = x.to(device)
with torch.no_grad():
if y_condition:
y = y.to(device)
z, nll, y_logits = model(x, y)
losses = compute_loss_y(nll,
y_logits,
y_weight,
y,
multi_class,
reduction='none')
else:
z, nll, y_logits = model(x, None)
losses = compute_loss(nll, reduction='none')
return losses
trainer = Engine(step)
checkpoint_handler = ModelCheckpoint(output_dir,
'glow',
save_interval=1,
n_saved=2,
require_empty=False)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler, {
'model': model,
'optimizer': optimizer
})
monitoring_metrics = ['total_loss']
RunningAverage(output_transform=lambda x: x['total_loss']).attach(
trainer, 'total_loss')
evaluator = Engine(eval_step)
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(lambda x, y: torch.mean(x),
output_transform=lambda x:
(x['total_loss'], torch.empty(x['total_loss'].shape[0]))).attach(
evaluator, 'total_loss')
if y_condition:
monitoring_metrics.extend(['nll'])
RunningAverage(output_transform=lambda x: x['nll']).attach(
trainer, 'nll')
# Note: replace by https://github.com/pytorch/ignite/pull/524 when released
Loss(lambda x, y: torch.mean(x),
output_transform=lambda x:
(x['nll'], torch.empty(x['nll'].shape[0]))).attach(
evaluator, 'nll')
pbar = ProgressBar()
pbar.attach(trainer, metric_names=monitoring_metrics)
# load pre-trained model if given
if saved_model:
model.load_state_dict(torch.load(saved_model))
model.set_actnorm_init()
if saved_optimizer:
optimizer.load_state_dict(torch.load(saved_optimizer))
file_name, ext = os.path.splitext(saved_model)
resume_epoch = int(file_name.split('_')[-1])
@trainer.on(Events.STARTED)
def resume_training(engine):
engine.state.epoch = resume_epoch
engine.state.iteration = resume_epoch * len(
engine.state.dataloader)
@trainer.on(Events.STARTED)
def init(engine):
model.train()
init_batches = []
init_targets = []
with torch.no_grad():
for batch, target in islice(train_loader, None, n_init_batches):
init_batches.append(batch)
init_targets.append(target)
init_batches = torch.cat(init_batches).to(device)
assert init_batches.shape[0] == n_init_batches * batch_size
if y_condition:
init_targets = torch.cat(init_targets).to(device)
else:
init_targets = None
model(init_batches, init_targets)
@trainer.on(Events.EPOCH_COMPLETED)
def evaluate(engine):
evaluator.run(test_loader)
metrics = evaluator.state.metrics
losses = ', '.join(
[f"{key}: {value:.2f}" for key, value in metrics.items()])
print(f'Validation Results - Epoch: {engine.state.epoch} {losses}')
timer = Timer(average=True)
timer.attach(trainer,
start=Events.EPOCH_STARTED,
resume=Events.ITERATION_STARTED,
pause=Events.ITERATION_COMPLETED,
step=Events.ITERATION_COMPLETED)
@trainer.on(Events.EPOCH_COMPLETED)
def print_times(engine):
pbar.log_message(
f'Epoch {engine.state.epoch} done. Time per batch: {timer.value():.3f}[s]'
)
timer.reset()
trainer.run(train_loader, epochs)
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
print(
ind,
args.delta,
log_p[i].item(),
logdet[i].item(),
train_labels[ind].item(),
file=f,
)
if ind >= 9999:
break
f.close()
if __name__ == "__main__":
args = parser.parse_args()
print(string_args(args))
device = args.device
model_single = Glow(
args.n_channels,
args.n_flow,
args.n_block,
affine=args.affine,
conv_lu=not args.no_lu,
)
model = model_single
model.load_state_dict(torch.load(args.model_path))
model = model.to(device)
test(args, model)
model_single.reverse(z_sample).cpu().data,
f'sample/{str(i + 1).zfill(6)}.png',
normalize=True,
nrow=10,
range=(-0.5, 0.5),
)
if i % 10000 == 0:
torch.save(model.state_dict(),
f'checkpoint/model_{str(i + 1).zfill(6)}.pt')
torch.save(optimizer.state_dict(),
f'checkpoint/optim_{str(i + 1).zfill(6)}.pt')
if __name__ == '__main__':
args = parser.parse_args()
print(args)
model_single = Glow(3,
args.n_flow,
args.n_block,
affine=args.affine,
conv_lu=not args.no_lu)
model = nn.DataParallel(model_single)
# model = model_single
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
train(args, model, optimizer)
with open(output_folder + 'hparams.json') as json_file:
hparams = json.load(json_file)
image_shape = (64, 64, 3)
num_classes = 40
Batch_Size = 4
dataset_test = CelebALoader(
root_folder=hparams['dataroot']
) #'/home/yellow/deep-learning-and-practice/hw7/dataset/task_2/'
test_loader = DataLoader(dataset_test,
batch_size=Batch_Size,
shuffle=False,
drop_last=True)
model = Glow(image_shape, hparams['hidden_channels'], hparams['K'],
hparams['L'], hparams['actnorm_scale'],
hparams['flow_permutation'], hparams['flow_coupling'],
hparams['LU_decomposed'], num_classes, hparams['learn_top'],
hparams['y_condition'])
model.load_state_dict(
torch.load(output_folder + model_name, map_location="cpu")['model'])
model.set_actnorm_init()
model = model.to(device)
model = model.eval()
# attribute_list = [8] # Black_Hair
attribute_list = [20, 31, 33] # Male, Smiling, Wavy_Hair, 24z No_Beard
# attribute_list = [11, 26, 31, 8, 6, 7] # Brown_Hair, Pale_Skin, Smiling, Black_Hair, Big_Lips, Big_Nose
# attribute_list = [i for i in range(40)]
N = 8
z_pos_list = [torch.Tensor([]).cuda() for i in range(len(attribute_list))]
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)
ldtv,
) in zip(log_p_val, logdet_val, log_p_train_val,
logdet_train_val):
print(
args.delta,
lpv.item(),
ldv.item(),
lptv.item(),
ldtv.item(),
file=f_ll,
)
f_ll.close()
f_train_loss.close()
f_test_loss.close()
if __name__ == "__main__":
args = parser.parse_args()
print(string_args(args))
device = args.device
model = Glow(
args.n_channels,
args.n_flow,
args.n_block,
affine=args.affine,
conv_lu=not args.no_lu,
)
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
train(args, model, optimizer)
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)
