def viz_forward_trajectory(data, forward_diffusion, scl, shft):
temperature = args.temperature
for num_step in range(args.num_steps):
x_recons, _ = forward_diffusion(data, temperature)
temperature = temperature * args.temperature_factor
x_recons = np.asarray(x_recons).astype('float32')
x_recons = x_recons.reshape(args.batch_size, INPUT_SIZE)
x_recons = x_recons.reshape(args.batch_size, 1, 28, 28)
x_recons = ((x_recons - shft) / scl)
plot_images(x_recons, 'forward_' + str(num_step))
def viz_forward_trajectory(data, forward_diffusion, scl, shft):
temperature = args.temperature
for num_step in range(args.num_steps):
x_recons, _ = forward_diffusion(data, temperature)
temperature = temperature * 2
x_recons = np.asarray(x_recons).astype('float32')
x_recons = x_recons.reshape(args.batch_size, INPUT_SIZE)
x_recons = x_recons.reshape(args.batch_size, N_COLORS, WIDTH, WIDTH)
x_recons = ((x_recons - shft) / scl)
plot_images(x_recons, str(num_step) + '_forward')
def generate_samples(model,
get_mu_sigma,
n_samples=36,
inpaint=False,
denoise_sigma=None,
X_true=None,
base_fname_part1="samples",
base_fname_part2='',
num_intermediate_plots=4,
seed=12345):
"""
Run the reverse diffusion process (generative model).
"""
# use the same noise in the samples every time, so they're easier to
# compare across learning
rng = np.random.RandomState(seed)
spatial_width = model.spatial_width
n_colors = model.n_colors
# set the initial state X^T of the reverse trajectory
XT = rng.normal(size=(n_samples, n_colors, spatial_width, spatial_width))
if denoise_sigma is not None:
XT = X_true + XT * denoise_sigma
base_fname_part1 += '_denoise%g' % denoise_sigma
if inpaint:
mask = generate_inpaint_mask(n_samples, n_colors, spatial_width)
XT.flat[mask] = X_true.flat[mask]
base_fname_part1 += '_inpaint'
else:
mask = None
if X_true is not None:
viz.plot_images(X_true, base_fname_part1 + '_true' + base_fname_part2)
viz.plot_images(
XT, base_fname_part1 + '_t%04d' % model.trajectory_length +
base_fname_part2)
Xmid = XT.copy()
for t in xrange(model.trajectory_length - 1, 0, -1):
Xmid = diffusion_step(Xmid, t, get_mu_sigma, denoise_sigma, mask, XT,
rng)
if np.mod(
model.trajectory_length - t,
int(
np.ceil(model.trajectory_length /
(num_intermediate_plots + 2.)))) == 0:
viz.plot_images(Xmid,
base_fname_part1 + '_t%04d' % t + base_fname_part2)
X0 = Xmid
viz.plot_images(X0, base_fname_part1 + '_t%04d' % 0 + base_fname_part2)
def generate_samples(model, get_mu_sigma, n_samples=36,
inpaint=False, denoise_sigma=None, logr_grad=None,
X_true=None,
base_fname_part1="samples", base_fname_part2='',
num_intermediate_plots=4, seed=12345):
"""
Run the reverse diffusion process (generative model).
"""
# use the same noise in the samples every time, so they're easier to
# compare across learning
rng = np.random.RandomState(seed)
spatial_width = model.spatial_width
n_colors = model.n_colors
# set the initial state X^T of the reverse trajectory
XT = rng.normal(size=(n_samples,n_colors,spatial_width,spatial_width))
if denoise_sigma is not None:
XT = X_true + XT*denoise_sigma
base_fname_part1 += '_denoise%g'%denoise_sigma
if inpaint:
mask = generate_inpaint_mask(n_samples, n_colors, spatial_width)
XT.flat[mask] = X_true.flat[mask]
base_fname_part1 += '_inpaint'
if logr_grad is not None:
base_fname_part1 += '_logrperturb'
else:
mask = None
if X_true is not None:
viz.plot_images(X_true, base_fname_part1 + '_true' + base_fname_part2)
viz.plot_images(XT, base_fname_part1 + '_t%04d'%model.trajectory_length + base_fname_part2)
Xmid = XT.copy()
for t in xrange(model.trajectory_length-1, 0, -1):
Xmid = diffusion_step(Xmid, t, get_mu_sigma, denoise_sigma, mask, XT, rng,
model.trajectory_length, logr_grad)
if np.mod(model.trajectory_length-t,
int(np.ceil(model.trajectory_length/(num_intermediate_plots+2.)))) == 0:
viz.plot_images(Xmid, base_fname_part1 + '_t%04d'%t + base_fname_part2)
X0 = Xmid
viz.plot_images(X0, base_fname_part1 + '_t%04d'%0 + base_fname_part2)
def train(args, model_args, lrate):
print("Copying the dataset to the current node's dir...")
tmp = '/Tmp/vermavik/'
home = '/u/vermavik/'
"""
tmp='/tmp/vermav1/'
home='/u/79/vermav1/unix/'
"""
dataset = args.dataset
data_source_dir = home + 'data/' + dataset + '/'
"""
if not os.path.exists(data_source_dir):
os.makedirs(data_source_dir)
data_target_dir = tmp+'data/CelebA/'
copy_tree(data_source_dir, data_target_dir)
"""
### set up the experiment directories########
exp_name = experiment_name(dataset=args.dataset,
act=args.activation,
meta_steps=args.meta_steps,
sigma=args.sigma,
temperature_factor=args.temperature_factor,
alpha1=args.alpha1,
alpha2=args.alpha2,
alpha3=args.alpha3,
grad_norm_max=args.grad_max_norm,
epochs=args.epochs,
job_id=args.job_id,
add_name=args.add_name)
#temp_model_dir = tmp+'experiments/HVWB/'+dataset+'/model/'+ exp_name
#temp_result_dir = tmp+'experiments/HVWB/'+dataset+'/results/'+ exp_name
model_dir = home + 'experiments/HVWB/' + dataset + '/model/' + exp_name
result_dir = home + 'experiments/HVWB/' + dataset + '/results/' + exp_name
if not os.path.exists(model_dir):
os.makedirs(model_dir)
#if not os.path.exists(temp_result_dir):
# os.makedirs(temp_result_dir)
"""
#copy_script_to_folder(os.path.abspath(__file__), temp_result_dir)
result_path = os.path.join(temp_result_dir , 'out.txt')
filep = open(result_path, 'w')
out_str = str(args)
print(out_str)
filep.write(out_str + '\n')
#torch.backends.cudnn.enabled = False # slower but repeatable
torch.backends.cudnn.enabled = True # faster but not repeatable
out_str = 'initial seed = ' + str(args.manualSeed)
print(out_str)
filep.write(out_str + '\n\n')
"""
#model_id = '/data/lisatmp4/anirudhg/minst_walk_back/walkback_'
"""
model_id = '/data/lisatmp4/anirudhg/celebA_latent_walkback/walkback_'
model_dir = create_log_dir(args, model_id)
model_id2 = '../celebA_logs/walkback_'
model_dir2 = create_log_dir(args, model_id2)
print model_dir
print model_dir2 + '/' + 'log.jsonl.gz'
logger = mimir.Logger(filename=model_dir2 + '/log.jsonl.gz', formatter=None)
"""
# TODO batches_per_epoch should not be hard coded
lrate = args.lr
import sys
sys.setrecursionlimit(10000000)
args, model_args = parse_args()
print args
## load the training data
print 'loading mnist'
train_loader, test_loader = load_mnist(
data_aug=0,
batch_size=100,
test_batch_size=100,
cuda=True,
data_target_dir="/u/vermavik/DARC/mnist")
#train_loader, unlabelled_loader, test_loader = get_mnist(location="/u/vermavik/DARC/mnist", batch_size=64, labels_per_class=100)
n_colors = 1
spatial_width = 28
for batch_idx, (data, target) in enumerate(train_loader):
Xbatch = data.numpy()
#print Xbatch
scl = 1. / np.sqrt(np.mean((Xbatch - np.mean(Xbatch))**2))
shft = -np.mean(Xbatch * scl)
break ### TO DO : calculate statistics on whole data
print "Width", WIDTH, spatial_width
model = Net(args)
if args.cuda:
model.cuda()
loss_fn = nn.BCELoss()
if args.optimizer == 'sgd':
optimizer_encoder = optim.SGD(model.encoder_params,
lr=args.lr,
momentum=args.momentum,
weight_decay=0)
optimizer_transition = optim.SGD(model.transition_params,
lr=args.lr,
momentum=args.momentum,
weight_decay=0)
optimizer_decoder = optim.SGD(model.decoder_params,
lr=args.lr,
momentum=args.momentum,
weight_decay=0)
elif args.optimizer == 'adam':
optimizer_encoder = optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0)
optimizer_transition = optim.Adam(model.transition_params,
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0)
optimizer_decoder = optim.Adam(model.decoder_params,
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0)
uidx = 0
estop = False
bad_counter = 0
#batch_index = 1
n_samples = 0
print 'Number of steps....'
print args.num_steps
print "Number of metasteps...."
print args.meta_steps
print 'Done'
count_sample = 1
#### for saving metrics for all steps ###
train_loss = []
train_x_loss = []
train_log_p_reverse = []
train_kld = []
#### for saving metrics for each step individually ###
train_loss_each_step = [[]]
train_x_loss_each_step = [[]]
train_log_p_reverse_each_step = [[]]
#train_kld_each_step = [[]]
for i in range(args.meta_steps - 1):
train_loss_each_step.append([])
train_x_loss_each_step.append([])
train_log_p_reverse_each_step.append([])
#train_kld_each_step.append([])
for epoch in range(args.epochs):
print('epoch', epoch)
for batch_idx, (data, target) in enumerate(train_loader):
#batch_idx = 0
#for (data, target), (u, _) in zip(cycle(train_loader), unlabelled_loader):
# batch_idx +=1
if args.cuda:
data, target = data.cuda(), target.cuda()
data, target = Variable(data), Variable(target)
data = data.view(-1, 1 * 28 * 28)
#print data.shape
t0 = time.time()
#batch_index += 1
n_samples += data.data.shape[0]
#print (n_samples)
temperature_forward = args.temperature
meta_cost = []
x = data
z = None
encode = True
for meta_step in range(0, args.meta_steps):
#print ('meta_step', meta_step)
#print encode
loss, x_loss, log_p_reverse, KLD, z, z_tilde, x_tilde = compute_loss(
x,
z,
model,
loss_fn,
temperature_forward,
meta_step,
encode=encode)
#meta_cost.append(loss)
#print compute_param_norm(model.conv_x_z_1.weight.data)
optimizer_encoder.zero_grad()
optimizer_transition.zero_grad()
optimizer_decoder.zero_grad()
loss.backward()
total_norm = clip_grad_norm(model.parameters(),
args.grad_max_norm)
#print ('step', meta_step, total_norm)
if encode == True:
optimizer_encoder.step()
optimizer_transition.step()
optimizer_decoder.step()
#print ('step', meta_step, clip_grad_norm(model.parameters(), 1000000))
### store metrics#######
train_loss.append(loss.data[0])
train_x_loss.append(x_loss.data[0])
train_log_p_reverse.append(-log_p_reverse.data[0])
if KLD is not None:
train_kld.append(KLD.data[0])
#### store metrices for each step separately###
train_loss_each_step[meta_step].append(loss.data[0])
train_x_loss_each_step[meta_step].append(x_loss.data[0])
train_log_p_reverse_each_step[meta_step].append(
-log_p_reverse.data[0])
#if KLD is not None:
# train_kld_each_step[meta_step].append(KLD.data[0])
if args.meta_steps > 1:
#data, _, _, _, _, _, _ = forward_diffusion(data, model, loss_fn,temperature_forward,meta_step)
#data = data.view(-1,3, 64,64)
#data = Variable(data.data, requires_grad=False)
x = Variable(x_tilde.data, requires_grad=False)
z = Variable(z_tilde.data, requires_grad=False)
if args.encode_every_step == 0:
encode = False
temperature_forward *= args.temperature_factor
#print loss.data
#print loss.data
#cost = sum(meta_cost) / len(meta_cost)
#print cost
#gradient_updates_ = get_grads(data_use[0],args.temperature)
if np.isnan(loss.data.cpu()[0]) or np.isinf(loss.data.cpu()[0]):
print loss.data
print 'NaN detected'
return 1.
#batch_idx=0
if batch_idx % 100 == 0:
plot_loss(model_dir, train_loss, train_x_loss,
train_log_p_reverse, train_kld, train_loss_each_step,
train_x_loss_each_step,
train_log_p_reverse_each_step, args.meta_steps)
count_sample += 1
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
temperature_forward = args.temperature
#print 'this'
data_forward_diffusion = data
for num_step in range(args.num_steps * args.meta_steps):
#print "Forward temperature", temperature_forward
data_forward_diffusion, _, _, _, _, _, _ = forward_diffusion(
data_forward_diffusion, model, loss_fn,
temperature_forward, num_step)
#print data_forward_diffusion.shape
#data_forward_diffusion = np.asarray(data).astype('float32').reshape(args.batch_size, INPUT_SIZE)
data_ = data_forward_diffusion.view(
-1, 1, spatial_width, spatial_width
) #reshape(args.batch_size, n_colors, WIDTH, WIDTH)
if num_step % 2 == 1:
plot_images(
data_.data.cpu().numpy(),
model_dir + '/' + "batch_" + str(batch_idx) +
'_corrupted_' + 'epoch_' + str(epoch) +
'_time_step_' + str(num_step))
temperature_forward = temperature_forward * args.temperature_factor
print "PLOTTING ORIGINAL IMAGE"
temp = data.view(-1, 1, spatial_width, spatial_width)
plot_images(
temp.data.cpu().numpy(), model_dir + '/' + 'orig_' +
'epoch_' + str(epoch) + '_batch_index_' + str(batch_idx))
print "DONE PLOTTING ORIGINAL IMAGE"
'''
temperature = args.temperature * (args.temperature_factor ** (args.num_steps*args.meta_steps - 1 ))
for i in range(args.num_steps*args.meta_steps + args.extra_steps):
x_data, sampled, sampled_activation, sampled_preactivation = f_sample(x_data, temperature)
print 'On backward step number, using temperature', i, temperature
reverse_time(scl, shft, x_data, model_dir + '/'+ "batch_" + str(batch_index) + '_samples_backward_' + 'epoch_' + str(count_sample) + '_time_step_' + str(i))
x_data = np.asarray(x_data).astype('float32')
x_data = x_data.reshape(args.batch_size, INPUT_SIZE)
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
'''
#print 'this'
if args.noise == "gaussian":
z_sampled = np.random.normal(
0.0, 1.0,
size=(args.batch_size, args.nl)) #.clip(0.0, 1.0)
else:
z_sampled = np.random.binomial(1,
0.5,
size=(args.batch_size,
args.nl))
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
z = torch.from_numpy(np.asarray(z_sampled).astype('float32'))
if args.cuda:
z = z.cuda()
z = Variable(z)
for i in range(args.num_steps *
args.meta_steps): # + args.extra_steps):
z_new_to_x, z_to_x, z_new = model.sample(
z, temperature,
args.num_steps * args.meta_steps - i - 1)
#print 'On step number, using temperature', i, temperature
if i % 2 == 1:
reverse_time(
scl, shft,
z_new_to_x.data.cpu().numpy(), model_dir +
'/batch_index_' + str(batch_idx) + '_inference_' +
'epoch_' + str(epoch) + '_step_' + str(i))
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
z = z_new
def reverse_time(scl, shft, sample_drawn, name):
new_image = np.asarray(sample_drawn).astype('float32').reshape(
args.batch_size, N_COLORS, WIDTH, WIDTH)
plot_images(new_image, name)
g_loss, x_gen, x_gen_corr = train_g(xs)
if iteration % 500 == 0:
print "=============================================="
print iteration
print "c value", c_params['log_sigma'].get_value().mean()
print "mean match orig", xs.mean(), x_gen.mean()
print "std match orig", xs.std(), x_gen.std()
print "mean match corr", x_corr.mean(), x_gen_corr.mean()
print "std match corr", x_corr.std(), x_gen_corr.std()
plot_images(xs.reshape((128, 1, 28, 28)),
"plots/" + slurm_name + "_real.png")
plot_images(x_gen.reshape((128, 1, 28, 28)),
"plots/" + slurm_name + "_gen.png")
plot_images(
x_corr.reshape((128, 1, 28, 28)).clip(0.0, 1.0),
"plots/" + slurm_name + "_corrupt_real")
plot_images(
x_gen_corr.reshape((128, 1, 28, 28)).clip(0.0, 1.0),
"plots/" + slurm_name + "_corrupt_gen")
if do_plot:
plt.hist(xs, alpha=0.5)
plt.hist(x_gen, alpha=0.5)
plt.legend(["real", "fake"])
plt.show()
def train(args, model_args):
#model_id = '/data/lisatmp4/lambalex/lsun_walkback/walkback_'
model_id = '/data/lisatmp4/anirudhg/cifar_walk_back/walkback_'
model_dir = create_log_dir(args, model_id)
model_id2 = 'logs/walkback_'
model_dir2 = create_log_dir(args, model_id2)
print model_dir
print model_dir2 + '/' + 'log.jsonl.gz'
logger = mimir.Logger(filename=model_dir2 + '/log.jsonl.gz',
formatter=None)
# TODO batches_per_epoch should not be hard coded
lrate = args.lr
import sys
sys.setrecursionlimit(10000000)
args, model_args = parse_args()
#trng = RandomStreams(1234)
if args.resume_file is not None:
print "Resuming training from " + args.resume_file
from blocks.scripts import continue_training
continue_training(args.resume_file)
## load the training data
if args.dataset == 'MNIST':
print 'loading MNIST'
from fuel.datasets import MNIST
dataset_train = MNIST(['train'], sources=('features', ))
dataset_test = MNIST(['test'], sources=('features', ))
n_colors = 1
spatial_width = 28
elif args.dataset == 'CIFAR10':
from fuel.datasets import CIFAR10
dataset_train = CIFAR10(['train'], sources=('features', ))
dataset_test = CIFAR10(['test'], sources=('features', ))
n_colors = 3
spatial_width = 32
elif args.dataset == "lsun" or args.dataset == "lsunsmall":
print "loading lsun class!"
from load_lsun import load_lsun
print "loading lsun data!"
if args.dataset == "lsunsmall":
dataset_train, dataset_test = load_lsun(args.batch_size,
downsample=True)
spatial_width = 32
else:
dataset_train, dataset_test = load_lsun(args.batch_size,
downsample=False)
spatial_width = 64
n_colors = 3
elif args.dataset == "celeba":
print "loading celeba data"
from fuel.datasets.celeba import CelebA
dataset_train = CelebA(which_sets=['train'],
which_format="64",
sources=('features', ),
load_in_memory=False)
dataset_test = CelebA(which_sets=['test'],
which_format="64",
sources=('features', ),
load_in_memory=False)
spatial_width = 64
n_colors = 3
tr_scheme = SequentialScheme(examples=dataset_train.num_examples,
batch_size=args.batch_size)
ts_scheme = SequentialScheme(examples=dataset_test.num_examples,
batch_size=args.batch_size)
train_stream = DataStream.default_stream(dataset_train,
iteration_scheme=tr_scheme)
test_stream = DataStream.default_stream(dataset_test,
iteration_scheme=ts_scheme)
dataset_train = train_stream
dataset_test = test_stream
#epoch_it = train_stream.get_epoch_iterator()
elif args.dataset == 'Spiral':
print 'loading SPIRAL'
train_set = Spiral(num_examples=100000,
classes=1,
cycles=2.,
noise=0.01,
sources=('features', ))
dataset_train = DataStream.default_stream(
train_set,
iteration_scheme=ShuffledScheme(train_set.num_examples,
args.batch_size))
else:
raise ValueError("Unknown dataset %s." % args.dataset)
model_options = locals().copy()
if args.dataset != 'lsun' and args.dataset != 'celeba':
train_stream = Flatten(
DataStream.default_stream(
dataset_train,
iteration_scheme=ShuffledScheme(
examples=dataset_train.num_examples -
(dataset_train.num_examples % args.batch_size),
batch_size=args.batch_size)))
else:
train_stream = dataset_train
test_stream = dataset_test
print "Width", WIDTH, spatial_width
shp = next(train_stream.get_epoch_iterator())[0].shape
print "got epoch iterator"
# make the training data 0 mean and variance 1
# TODO compute mean and variance on full dataset, not minibatch
Xbatch = next(train_stream.get_epoch_iterator())[0]
scl = 1. / np.sqrt(np.mean((Xbatch - np.mean(Xbatch))**2))
shft = -np.mean(Xbatch * scl)
# scale is applied before shift
#train_stream = ScaleAndShift(train_stream, scl, shft)
#test_stream = ScaleAndShift(test_stream, scl, shft)
print 'Building model'
params = init_params(model_options)
if args.reload_:
print "Trying to reload parameters"
if os.path.exists(args.saveto_filename):
print 'Reloading Parameters'
print args.saveto_filename
params = load_params(args.saveto_filename, params)
tparams = init_tparams(params)
print tparams
'''
x = T.matrix('x', dtype='float32')
temp = T.scalar('temp', dtype='float32')
f=transition_operator(tparams, model_options, x, temp)
for data in train_stream.get_epoch_iterator():
print data[0]
a = f([data[0], 1.0, 1])
#ipdb.set_trace()
'''
x, cost, start_temperature = build_model(tparams, model_options)
inps = [x, start_temperature]
x_Data = T.matrix('x_Data', dtype='float32')
temperature = T.scalar('temperature', dtype='float32')
forward_diffusion = one_step_diffusion(x_Data, model_options, tparams,
temperature)
#print 'Building f_cost...',
#f_cost = theano.function(inps, cost)
#print 'Done'
print tparams
grads = T.grad(cost, wrt=itemlist(tparams))
#get_grads = theano.function(inps, grads)
for j in range(0, len(grads)):
grads[j] = T.switch(T.isnan(grads[j]), T.zeros_like(grads[j]),
grads[j])
# compile the optimizer, the actual computational graph is compiled here
lr = T.scalar(name='lr')
print 'Building optimizers...',
optimizer = args.optimizer
f_grad_shared, f_update = getattr(optimizers, optimizer)(lr, tparams,
grads, inps, cost)
print 'Done'
for param in tparams:
print param
print tparams[param].get_value().shape
print 'Buiding Sampler....'
f_sample = sample(tparams, model_options)
print 'Done'
uidx = 0
estop = False
bad_counter = 0
max_epochs = 4000
batch_index = 1
print 'Number of steps....'
print args.num_steps
print "Number of metasteps...."
print args.meta_steps
print 'Done'
count_sample = 1
for eidx in xrange(max_epochs):
if eidx % 20 == 0:
params = unzip(tparams)
save_params(params,
model_dir + '/' + 'params_' + str(eidx) + '.npz')
n_samples = 0
print 'Starting Next Epoch ', eidx
for data in train_stream.get_epoch_iterator():
if args.dataset == 'CIFAR10':
if data[0].shape[0] == args.batch_size:
data_use = (data[0].reshape(args.batch_size,
3 * 32 * 32), )
else:
continue
t0 = time.time()
batch_index += 1
n_samples += len(data_use[0])
uidx += 1
if data_use[0] is None:
print 'No data '
uidx -= 1
continue
ud_start = time.time()
t1 = time.time()
data_run = data_use[0]
temperature_forward = args.temperature
meta_cost = []
for meta_step in range(0, args.meta_steps):
meta_cost.append(f_grad_shared(data_run, temperature_forward))
f_update(lrate)
if args.meta_steps > 1:
data_run, sigma, _, _ = forward_diffusion(
[data_run, temperature_forward, 1])
temperature_forward *= args.temperature_factor
cost = sum(meta_cost) / len(meta_cost)
ud = time.time() - ud_start
#gradient_updates_ = get_grads(data_use[0],args.temperature)
if np.isnan(cost) or np.isinf(cost):
print 'NaN detected'
return 1.
t1 = time.time()
#print time.time() - t1, "time to get grads"
t1 = time.time()
logger.log({
'epoch': eidx,
'batch_index': batch_index,
'uidx': uidx,
'training_error': cost
})
#'Norm_1': np.linalg.norm(gradient_updates_[0]),
#'Norm_2': np.linalg.norm(gradient_updates_[1]),
#'Norm_3': np.linalg.norm(gradient_updates_[2]),
#'Norm_4': np.linalg.norm(gradient_updates_[3])})
#print time.time() - t1, "time to log"
#print time.time() - t0, "total time in batch"
t5 = time.time()
if batch_index % 20 == 0:
print batch_index, "cost", cost
if batch_index % 200 == 0:
count_sample += 1
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
temperature_forward = args.temperature
for num_step in range(args.num_steps * args.meta_steps):
print "Forward temperature", temperature_forward
if num_step == 0:
x_data, sampled, sampled_activation, sampled_preactivation = forward_diffusion(
[data_use[0], temperature_forward, 1])
x_data = np.asarray(x_data).astype('float32').reshape(
args.batch_size, INPUT_SIZE)
x_temp = x_data.reshape(args.batch_size, n_colors,
WIDTH, WIDTH)
plot_images(
x_temp, model_dir + '/' + "batch_" +
str(batch_index) + '_corrupted' + 'epoch_' +
str(count_sample) + '_time_step_' + str(num_step))
else:
x_data, sampled, sampled_activation, sampled_preactivation = forward_diffusion(
[x_data, temperature_forward, 1])
x_data = np.asarray(x_data).astype('float32').reshape(
args.batch_size, INPUT_SIZE)
x_temp = x_data.reshape(args.batch_size, n_colors,
WIDTH, WIDTH)
plot_images(
x_temp, model_dir + '/batch_' + str(batch_index) +
'_corrupted' + '_epoch_' + str(count_sample) +
'_time_step_' + str(num_step))
temperature_forward = temperature_forward * args.temperature_factor
x_temp2 = data_use[0].reshape(args.batch_size, n_colors, WIDTH,
WIDTH)
plot_images(
x_temp2, model_dir + '/' + 'orig_' + 'epoch_' + str(eidx) +
'_batch_index_' + str(batch_index))
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
for i in range(args.num_steps * args.meta_steps +
args.extra_steps):
x_data, sampled, sampled_activation, sampled_preactivation = f_sample(
[x_data, temperature, 0])
print 'On backward step number, using temperature', i, temperature
reverse_time(
scl, shft, x_data, model_dir + '/' + "batch_" +
str(batch_index) + '_samples_backward_' + 'epoch_' +
str(count_sample) + '_time_step_' + str(i))
x_data = np.asarray(x_data).astype('float32')
x_data = x_data.reshape(args.batch_size, INPUT_SIZE)
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
if args.noise == "gaussian":
x_sampled = np.random.normal(
0.5, 2.0,
size=(args.batch_size, INPUT_SIZE)).clip(0.0, 1.0)
else:
s = np.random.binomial(1, 0.5, INPUT_SIZE)
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
x_data = np.asarray(x_sampled).astype('float32')
for i in range(args.num_steps * args.meta_steps +
args.extra_steps):
x_data, sampled, sampled_activation, sampled_preactivation = f_sample(
[x_data, temperature, 0])
print 'On step number, using temperature', i, temperature
reverse_time(
scl, shft, x_data, model_dir + '/batch_index_' +
str(batch_index) + '_inference_' + 'epoch_' +
str(count_sample) + '_step_' + str(i))
x_data = np.asarray(x_data).astype('float32')
x_data = x_data.reshape(args.batch_size, INPUT_SIZE)
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
ipdb.set_trace()
#print 'this'
data_forward_diffusion = data
for num_step in range(args.num_steps * args.meta_steps):
<<<<<<< HEAD
print "Forward temperature", temperature_forward
=======
#print "Forward temperature", temperature_forward
>>>>>>> 3cecdbb5646fed5f0ebca8204c9149e627a4ad2d
data_forward_diffusion, _, _, _, _, _, _ = forward_diffusion(data_forward_diffusion, model, loss_fn,temperature_forward, num_step)
#print data_forward_diffusion.shape
#data_forward_diffusion = np.asarray(data).astype('float32').reshape(args.batch_size, INPUT_SIZE)
data_forward_diffusion = data_forward_diffusion.view(-1,3, 64,64)#reshape(args.batch_size, n_colors, WIDTH, WIDTH)
<<<<<<< HEAD
plot_images(data_forward_diffusion.data.cpu().numpy(), model_dir + '/' + "batch_" + str(batch_idx) + '_corrupted_' + 'epoch_' + str(epoch) + '_time_step_' + str(num_step))
=======
if num_step%2==1:
plot_images(data_forward_diffusion.data.cpu().numpy(), model_dir + '/' + "batch_" + str(batch_idx) + '_corrupted_' + 'epoch_' + str(epoch) + '_time_step_' + str(num_step))
>>>>>>> 3cecdbb5646fed5f0ebca8204c9149e627a4ad2d
temperature_forward = temperature_forward * args.temperature_factor;
print "PLOTTING ORIGINAL IMAGE"
temp = data
plot_images(temp.data.cpu().numpy() , model_dir + '/' + 'orig_' + 'epoch_' + str(epoch) + '_batch_index_' + str(batch_idx))
print "DONE PLOTTING ORIGINAL IMAGE"
'''
xu = sample_data()
if iteration % 100 == 0:
print "===================================="
for step in range(0, num_steps):
xq = rng.uniform(-1.5, 1.5, size=(16, 28 * 28)).astype('float32')
out = train(xu, xq)
xu = out['xn']
s_xl = out['s_xl']
if iteration % 100 == 0:
print step, "s", s_xl[0]
if step == num_steps - 1:
print "s for prior", get_s(xq)
print "xu shape", xu.shape
plot_images(xu.reshape(16, 1, 28, 28),
"plots/imagewalk_step_" + str(step))
if iteration % 2000 == 0:
print "DOING GENERATION"
for step in range(0, num_steps):
res = reverse_step(xu)
print step, "score", res[1][0]
xu = res[0]
plot_images(xu.reshape(16, 1, 28, 28),
"plots/gen/step_" + str(step))
def reverse_time(scl, shft, sample_drawn, name):
#new_image = ((sample_drawn-shft)/scl)
#new_image = new_image.reshape(args.batch-size, 1, 28, 28)
new_image = np.asarray(sample_drawn).astype('float32').reshape(
args.batch_size, 1, 28, 28)
plot_images(new_image, name)
def train(args, model_args, lrate):
print("Copying the dataset to the current node's dir...")
tmp = '/Tmp/vermavik/'
home = '/u/vermavik/'
"""
tmp='/tmp/vermav1/'
home='/u/79/vermav1/unix/'
"""
dataset = 'celebA'
data_source_dir = home + 'data/CelebAsmall/'
"""
if not os.path.exists(data_source_dir):
os.makedirs(data_source_dir)
data_target_dir = tmp+'data/CelebA/'
copy_tree(data_source_dir, data_target_dir)
"""
### set up the experiment directories########
"""
exp_name=experiment_name(args.epochs,
args.batch_size,
args.test_batch_size,
args.lr,
args.momentum,
args.alpha1,
args.alpha2,
args.alpha3,
args.data_aug,
args.job_id,
args.add_name)
"""
exp_name = 'temp'
temp_model_dir = tmp + 'experiments/HVWB/' + dataset + '/model/' + exp_name
temp_result_dir = tmp + 'experiments/HVWB/' + dataset + '/results/' + exp_name
model_dir = home + 'experiments/HVWB/' + dataset + '/model/' + exp_name
result_dir = home + 'experiments/HVWB/' + dataset + '/results/' + exp_name
if not os.path.exists(temp_model_dir):
os.makedirs(temp_model_dir)
if not os.path.exists(temp_result_dir):
os.makedirs(temp_result_dir)
"""
#copy_script_to_folder(os.path.abspath(__file__), temp_result_dir)
result_path = os.path.join(temp_result_dir , 'out.txt')
filep = open(result_path, 'w')
out_str = str(args)
print(out_str)
filep.write(out_str + '\n')
#torch.backends.cudnn.enabled = False # slower but repeatable
torch.backends.cudnn.enabled = True # faster but not repeatable
out_str = 'initial seed = ' + str(args.manualSeed)
print(out_str)
filep.write(out_str + '\n\n')
"""
#model_id = '/data/lisatmp4/anirudhg/minst_walk_back/walkback_'
"""
model_id = '/data/lisatmp4/anirudhg/celebA_latent_walkback/walkback_'
model_dir = create_log_dir(args, model_id)
model_id2 = '../celebA_logs/walkback_'
model_dir2 = create_log_dir(args, model_id2)
print model_dir
print model_dir2 + '/' + 'log.jsonl.gz'
logger = mimir.Logger(filename=model_dir2 + '/log.jsonl.gz', formatter=None)
"""
# TODO batches_per_epoch should not be hard coded
lrate = args.lr
import sys
sys.setrecursionlimit(10000000)
args, model_args = parse_args()
print args
## load the training data
print 'loading celebA'
train_loader, test_loader = load_celebA(args.data_aug, args.batch_size,
args.batch_size, args.cuda,
data_source_dir)
n_colors = 3
spatial_width = 64
for batch_idx, (data, target) in enumerate(train_loader):
Xbatch = data.numpy()
#print Xbatch
scl = 1. / np.sqrt(np.mean((Xbatch - np.mean(Xbatch))**2))
shft = -np.mean(Xbatch * scl)
break ### TO DO : calculate statistics on whole data
print "Width", WIDTH, spatial_width
model = Net(args)
if args.cuda:
model.cuda()
loss_fn = nn.BCELoss()
if args.optimizer == 'sgd':
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=0)
elif args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0)
uidx = 0
estop = False
bad_counter = 0
max_epochs = 100
batch_index = 1
n_samples = 0
print 'Number of steps....'
print args.num_steps
print "Number of metasteps...."
print args.meta_steps
print 'Done'
count_sample = 1
train_loss = []
train_x_loss = []
train_log_p_reverse = []
for epoch in range(max_epochs):
for batch_idx, (data, target) in enumerate(train_loader):
#for batch_idx in xrange(100):
# data = torch.randn(100,3,64,64)
# target = torch.randn(100,1)
#print (batch_idx)
if args.cuda:
data, target = data.cuda(), target.cuda()
data, target = Variable(data), Variable(target)
t0 = time.time()
#batch_index += 1
n_samples += data.data.shape[0]
#print (n_samples)
temperature_forward = args.temperature
meta_cost = []
optimizer.zero_grad()
for meta_step in range(0, args.meta_steps):
#print ('meta_step', meta_step)
loss, x_loss, log_p_reverse = compute_loss(
data, model, loss_fn, temperature_forward, meta_step)
#meta_cost.append(loss)
#print loss
loss.backward()
print('epoch=', epoch, 'batch_id=', batch_idx, 'meta_step=',
meta_step, 'norm=', compute_norm(model))
train_loss.append(loss.data[0])
train_x_loss.append(x_loss.data[0])
train_log_p_reverse.append(log_p_reverse.data[0])
if args.meta_steps > 1:
data, _, _, _, _, _, _ = forward_diffusion(
data, model, loss_fn, temperature_forward, meta_step)
data = data.view(-1, 3, 64, 64)
data = Variable(data.data, requires_grad=False)
temperature_forward *= args.temperature_factor
optimizer.step()
#print loss.data
#print loss.data
#cost = sum(meta_cost) / len(meta_cost)
#print cost
#gradient_updates_ = get_grads(data_use[0],args.temperature)
if np.isnan(loss.data.cpu()[0]) or np.isinf(loss.data.cpu()[0]):
print loss.data
print 'NaN detected'
return 1.
#batch_idx=0
if batch_idx % 100 == 0:
plot_loss(model_dir, train_loss, train_x_loss,
train_log_p_reverse)
count_sample += 1
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
temperature_forward = args.temperature
#print 'this'
data_forward_diffusion = data
for num_step in range(args.num_steps * args.meta_steps):
print "Forward temperature", temperature_forward
data_forward_diffusion, _, _, _, _, _, _ = forward_diffusion(
data_forward_diffusion, model, loss_fn,
temperature_forward, num_step)
#print data_forward_diffusion.shape
#data_forward_diffusion = np.asarray(data).astype('float32').reshape(args.batch_size, INPUT_SIZE)
data_forward_diffusion = data_forward_diffusion.view(
-1, 3, 64,
64) #reshape(args.batch_size, n_colors, WIDTH, WIDTH)
plot_images(
data_forward_diffusion.data.cpu().numpy(), model_dir +
'/' + "batch_" + str(batch_idx) + '_corrupted_' +
'epoch_' + str(epoch) + '_time_step_' + str(num_step))
temperature_forward = temperature_forward * args.temperature_factor
print "PLOTTING ORIGINAL IMAGE"
temp = data
plot_images(
temp.data.cpu().numpy(), model_dir + '/' + 'orig_' +
'epoch_' + str(epoch) + '_batch_index_' + str(batch_idx))
print "DONE PLOTTING ORIGINAL IMAGE"
'''
temperature = args.temperature * (args.temperature_factor ** (args.num_steps*args.meta_steps - 1 ))
for i in range(args.num_steps*args.meta_steps + args.extra_steps):
x_data, sampled, sampled_activation, sampled_preactivation = f_sample(x_data, temperature)
print 'On backward step number, using temperature', i, temperature
reverse_time(scl, shft, x_data, model_dir + '/'+ "batch_" + str(batch_index) + '_samples_backward_' + 'epoch_' + str(count_sample) + '_time_step_' + str(i))
x_data = np.asarray(x_data).astype('float32')
x_data = x_data.reshape(args.batch_size, INPUT_SIZE)
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
'''
#print 'this'
if args.noise == "gaussian":
z_sampled = np.random.normal(
0.0, 1.0,
size=(args.batch_size, args.nl)) #.clip(0.0, 1.0)
else:
z_sampled = np.random.binomial(1,
0.5,
size=(args.batch_size,
args.nl))
temperature = args.temperature * (args.temperature_factor**(
args.num_steps * args.meta_steps - 1))
z = torch.from_numpy(np.asarray(z_sampled).astype('float32'))
if args.cuda:
z = z.cuda()
z = Variable(z)
for i in range(args.num_steps *
args.meta_steps): # + args.extra_steps):
z_new_to_x, z_to_x, z_new = model.sample(
z, temperature,
args.num_steps * args.meta_steps - i - 1)
#print 'On step number, using temperature', i, temperature
reverse_time(
scl, shft,
z_new_to_x.data.cpu().numpy(), model_dir +
'/batch_index_' + str(batch_idx) + '_inference_' +
'epoch_' + str(epoch) + '_step_' + str(i))
if temperature == args.temperature:
temperature = temperature
else:
temperature /= args.temperature_factor
z = z_new
copy_tree(temp_model_dir, model_dir)
copy_tree(temp_result_dir, result_dir)
rmtree(temp_model_dir)
rmtree(temp_result_dir)
