def main(lr, batch_size, image_size, K, T, num_iter, gpu):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu[0])
data_path = "../data/vimeo_interp_test/target/"
f = open(data_path + "tri_testlist.txt", "r")
trainfiles = [l[:-1] for l in f.readlines()]
margin = 0.3
updateD = False
updateG = True
iters = 0
prefix = ("VIMEO_MCNET" + "_image_size=" + str(image_size) + "_K=" +
str(K) + "_T=" + str(T) + "_batch_size=" + str(batch_size) +
"_lr=" + str(lr))
print("\n" + prefix + "\n")
checkpoint_dir = "../models/" + prefix + "/"
samples_dir = "../samples/" + prefix + "/"
summary_dir = "../logs/" + prefix + "/"
if not exists(checkpoint_dir):
makedirs(checkpoint_dir)
if not exists(samples_dir):
makedirs(samples_dir)
if not exists(summary_dir):
makedirs(summary_dir)
with tf.device("/gpu:%d" % gpu[0]):
model = MCNET(image_size=[image_size, image_size],
c_dim=3,
K=K,
batch_size=batch_size,
T=T,
checkpoint_dir=checkpoint_dir)
g_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
model.L_img, var_list=model.g_vars)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)) as sess:
tf.global_variables_initializer().run()
if model.load(sess, checkpoint_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
g_sum = tf.summary.merge(
[model.L_p_sum, model.L_gdl_sum, model.loss_sum])
writer = tf.summary.FileWriter(summary_dir, sess.graph)
counter = iters + 1
start_time = time.time()
while iters < num_iter:
mini_batches = get_minibatches_idx(len(trainfiles),
batch_size,
shuffle=True)
for _, batchidx in mini_batches:
if len(batchidx) == batch_size:
seq_batch = np.zeros(
(batch_size, image_size, image_size, K + T, 3),
dtype="float32")
diff_batch = np.zeros(
(batch_size, image_size, image_size, K - 1, 3),
dtype="float32")
t0 = time.time()
Ts = np.repeat(np.array([T]), batch_size, axis=0)
Ks = np.repeat(np.array([K]), batch_size, axis=0)
paths = np.repeat(data_path, batch_size, axis=0)
tfiles = np.array(trainfiles)[batchidx]
shapes = np.repeat(np.array([image_size]),
batch_size,
axis=0)
output = [
load_vimeo_data(f, p, img_sze, k,
t) for f, p, img_sze, k, t in zip(
tfiles, paths, shapes, Ks, Ts)
]
for i in range(batch_size):
seq_batch[i] = output[i][0]
diff_batch[i] = output[i][1]
_, summary_str = sess.run(
[g_optim, g_sum],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
writer.add_summary(summary_str, counter)
errL_img = model.L_img.eval({
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
counter += 1
if counter % 50 == 0:
print("Iters: [%2d] time: %4.4f, L_img: %.8f" %
(iters, time.time() - start_time, errL_img))
if np.mod(counter, 10) == 1:
samples = sess.run(
[model.G],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})[0]
# IPython.embed()
samples = samples[0].swapaxes(0, 2).swapaxes(1, 2)
# IPython.embed()
sbatch = seq_batch[0, :, :, :].swapaxes(0,
2).swapaxes(1, 2)
sbatch2 = sbatch.copy()
# IPython.embed()
sbatch2[K:, :, :] = samples
# IPython.embed()
samples = np.concatenate((sbatch2, sbatch), axis=0)
# IPython.embed()
print("Saving sample ...")
save_images(samples, [2, K + T],
samples_dir + "train_%s.png" % (iters))
if np.mod(counter, 10000) == 2:
model.save(sess, checkpoint_dir, counter)
iters += 1
def main(lr, batch_size, alpha, beta, image_size, K, T, num_iter, gpu):
data_path = "../data/KTH/"
f = open(data_path + "train_data_list_trimmed.txt", "r")
trainfiles = f.readlines()
margin = 0.3
updateD = True
updateG = True
iters = 0
prefix = ("KTH_MCNET" + "_image_size=" + str(image_size) + "_K=" + str(K) +
"_T=" + str(T) + "_batch_size=" + str(batch_size) + "_alpha=" +
str(alpha) + "_beta=" + str(beta) + "_lr=" + str(lr))
print("\n" + prefix + "\n")
checkpoint_dir = "../models/" + prefix + "/"
samples_dir = "../samples/" + prefix + "/"
summary_dir = "../logs/" + prefix + "/"
if not exists(checkpoint_dir):
makedirs(checkpoint_dir)
if not exists(samples_dir):
makedirs(samples_dir)
if not exists(summary_dir):
makedirs(summary_dir)
with tf.device("/gpu:%d" % gpu[0]):
model = MCNET(image_size=[image_size, image_size],
c_dim=1,
K=K,
batch_size=batch_size,
T=T,
checkpoint_dir=checkpoint_dir)
d_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
model.d_loss, var_list=model.d_vars)
g_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
alpha * model.L_img + beta * model.L_GAN, var_list=model.g_vars)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)) as sess:
tf.global_variables_initializer().run()
if model.load(sess, checkpoint_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
g_sum = tf.summary.merge(
[model.L_p_sum, model.L_gdl_sum, model.loss_sum, model.L_GAN_sum])
d_sum = tf.summary.merge(
[model.d_loss_real_sum, model.d_loss_sum, model.d_loss_fake_sum])
writer = tf.summary.FileWriter(summary_dir, sess.graph)
counter = iters + 1
start_time = time.time()
with Parallel(n_jobs=batch_size) as parallel:
while iters < num_iter:
mini_batches = get_minibatches_idx(len(trainfiles),
batch_size,
shuffle=True)
for _, batchidx in mini_batches:
if len(batchidx) == batch_size:
seq_batch = np.zeros(
(batch_size, image_size, image_size, K + T, 1),
dtype="float32")
diff_batch = np.zeros(
(batch_size, image_size, image_size, K - 1, 1),
dtype="float32")
t0 = time.time()
Ts = np.repeat(np.array([T]), batch_size, axis=0)
Ks = np.repeat(np.array([K]), batch_size, axis=0)
paths = np.repeat(data_path, batch_size, axis=0)
tfiles = np.array(trainfiles)[batchidx]
shapes = np.repeat(np.array([image_size]),
batch_size,
axis=0)
output = parallel(
delayed(load_kth_data)(f, p, img_sze, k, t)
for f, p, img_sze, k, t in zip(
tfiles, paths, shapes, Ks, Ts))
for i in xrange(batch_size):
seq_batch[i] = output[i][0]
diff_batch[i] = output[i][1]
if updateD:
_, summary_str = sess.run(
[d_optim, d_sum],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
writer.add_summary(summary_str, counter)
if updateG:
_, summary_str = sess.run(
[g_optim, g_sum],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
writer.add_summary(summary_str, counter)
errD_fake = model.d_loss_fake.eval({
model.diff_in:
diff_batch,
model.xt:
seq_batch[:, :, :, K - 1],
model.target:
seq_batch
})
errD_real = model.d_loss_real.eval({
model.diff_in:
diff_batch,
model.xt:
seq_batch[:, :, :, K - 1],
model.target:
seq_batch
})
errG = model.L_GAN.eval({
model.diff_in:
diff_batch,
model.xt:
seq_batch[:, :, :, K - 1],
model.target:
seq_batch
})
if errD_fake < margin or errD_real < margin:
updateD = False
if errD_fake > (1. - margin) or errD_real > (1. -
margin):
updateG = False
if not updateD and not updateG:
updateD = True
updateG = True
counter += 1
print(
"Iters: [%2d] time: %4.4f, d_loss: %.8f, L_GAN: %.8f"
% (iters, time.time() - start_time,
errD_fake + errD_real, errG))
if np.mod(counter, 100) == 1:
samples = sess.run(
[model.G],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})[0]
samples = samples[0].swapaxes(0, 2).swapaxes(1, 2)
sbatch = seq_batch[0, :, :,
K:].swapaxes(0,
2).swapaxes(1, 2)
samples = np.concatenate((samples, sbatch), axis=0)
print("Saving sample ...")
save_images(samples[:, :, :, ::-1], [2, T],
samples_dir + "train_%s.png" % (iters))
if np.mod(counter, 500) == 2:
model.save(sess, checkpoint_dir, counter)
iters += 1
def main(prefix, image_size, K, T, gpu):
data_path = "../data/KTH/"
f = open(data_path+"test_data_list.txt","r")
testfiles = f.readlines()
c_dim = 1
iters = 0
if prefix == "paper_models":
checkpoint_dir = "../models/"+prefix+"/KTH/"
best_model = "MCNET.model-98502"
else:
checkpoint_dir = "../models/"+prefix+"/"
best_model = None # will pick last model
with tf.device("/gpu:%d"%gpu[0]):
model = MCNET(image_size=[image_size, image_size], batch_size=1, K=K,
T=T, c_dim=c_dim, checkpoint_dir=checkpoint_dir,
is_train=False)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)) as sess:
tf.global_variables_initializer().run()
loaded, model_name = model.load(sess, checkpoint_dir, best_model)
if loaded:
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed... exitting")
return
quant_dir = "../results/quantitative/KTH/"+prefix+"/"
save_path = quant_dir+"results_model="+model_name+".npz"
if not exists(quant_dir):
makedirs(quant_dir)
vid_names = []
psnr_err = np.zeros((0, T))
ssim_err = np.zeros((0, T))
for i in xrange(len(testfiles)):
tokens = testfiles[i].split()
vid_path = data_path+tokens[0]+"_uncomp.avi"
while True:
try:
vid = imageio.get_reader(vid_path,"ffmpeg")
break
except Exception:
print("imageio failed loading frames, retrying")
action = vid_path.split("_")[1]
if action in ["running", "jogging"]:
n_skip = 3
else:
n_skip = T
for j in xrange(int(tokens[1]),int(tokens[2])-K-T-1,n_skip):
print("Video "+str(i)+"/"+str(len(testfiles))+". Index "+str(j)+
"/"+str(vid.get_length()-T-1))
folder_pref = vid_path.split("/")[-1].split(".")[0]
folder_name = folder_pref+"."+str(j)+"-"+str(j+T)
vid_names.append(folder_name)
savedir = "../results/images/KTH/"+prefix+"/"+folder_name
seq_batch = np.zeros((1, image_size, image_size,
K+T, c_dim), dtype="float32")
diff_batch = np.zeros((1, image_size, image_size,
K-1, 1), dtype="float32")
for t in xrange(K+T):
# imageio fails randomly sometimes
while True:
try:
img = cv2.resize(vid.get_data(j+t), (image_size, image_size))
break
except Exception:
print("imageio failed loading frames, retrying")
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
seq_batch[0,:,:,t] = transform(img[:,:,None])
for t in xrange(1,K):
prev = inverse_transform(seq_batch[0,:,:,t-1])
next = inverse_transform(seq_batch[0,:,:,t])
diff = next.astype("float32")-prev.astype("float32")
diff_batch[0,:,:,t-1] = diff
true_data = seq_batch[:,:,:,K:,:].copy()
pred_data = np.zeros(true_data.shape, dtype="float32")
xt = seq_batch[:,:,:,K-1]
pred_data[0] = sess.run(model.G,
feed_dict={model.diff_in: diff_batch,
model.xt: xt})
if not os.path.exists(savedir):
os.makedirs(savedir)
cpsnr = np.zeros((K+T,))
cssim = np.zeros((K+T,))
pred_data = np.concatenate((seq_batch[:,:,:,:K], pred_data),axis=3)
true_data = np.concatenate((seq_batch[:,:,:,:K], true_data),axis=3)
for t in xrange(K+T):
pred = (inverse_transform(pred_data[0,:,:,t])*255).astype("uint8")
target = (inverse_transform(true_data[0,:,:,t])*255).astype("uint8")
cpsnr[t] = measure.compare_psnr(pred,target)
cssim[t] = ssim.compute_ssim(Image.fromarray(cv2.cvtColor(target,
cv2.COLOR_GRAY2BGR)),
Image.fromarray(cv2.cvtColor(pred,
cv2.COLOR_GRAY2BGR)))
pred = draw_frame(pred, t < K)
target = draw_frame(target, t < K)
cv2.imwrite(savedir+"/pred_"+"{0:04d}".format(t)+".png", pred)
cv2.imwrite(savedir+"/gt_"+"{0:04d}".format(t)+".png", target)
cmd1 = "rm "+savedir+"/pred.gif"
cmd2 = ("ffmpeg -f image2 -framerate 7 -i "+savedir+
"/pred_%04d.png "+savedir+"/pred.gif")
cmd3 = "rm "+savedir+"/pred*.png"
# Comment out "system(cmd3)" if you want to keep the output images
# Otherwise only the gifs will be kept
system(cmd1); system(cmd2); system(cmd3)
cmd1 = "rm "+savedir+"/gt.gif"
cmd2 = ("ffmpeg -f image2 -framerate 7 -i "+savedir+
"/gt_%04d.png "+savedir+"/gt.gif")
cmd3 = "rm "+savedir+"/gt*.png"
# Comment out "system(cmd3)" if you want to keep the output images
# Otherwise only the gifs will be kept
system(cmd1); system(cmd2); system(cmd3)
psnr_err = np.concatenate((psnr_err, cpsnr[None,K:]), axis=0)
ssim_err = np.concatenate((ssim_err, cssim[None,K:]), axis=0)
np.savez(save_path, psnr=psnr_err, ssim=ssim_err)
print("Results saved to "+save_path)
print("Done.")
def main(lr, prefix, K, T, gpu):
data_path = "../data/UCF101/UCF-101/"
f = open(data_path.rsplit("/", 2)[0] + "/testlist01.txt", "r")
testfiles = f.readlines()
image_size = [240, 320]
c_dim = 3
iters = 0
if prefix == "paper_models":
checkpoint_dir = "../models/" + prefix + "/S1M/"
best_model = "MCNET.model-102502"
else:
checkpoint_dir = "../models/" + prefix + "/"
best_model = None # will pick last model
with tf.device("/gpu:%d" % gpu[0]):
model = MCNET(image_size=image_size,
batch_size=1,
K=K,
T=T,
c_dim=c_dim,
checkpoint_dir=checkpoint_dir,
is_train=False)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)) as sess:
tf.global_variables_initializer().run()
loaded, model_name = model.load(sess, checkpoint_dir, best_model)
if loaded:
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed... exitting")
return
quant_dir = "../results/quantitative/UCF101/" + prefix + "/"
save_path = quant_dir + "results_model=" + model_name + ".npz"
if not exists(quant_dir):
makedirs(quant_dir)
vid_names = []
psnr_err = np.zeros((0, T))
ssim_err = np.zeros((0, T))
for i in range(0, len(testfiles), 10):
print(" Video " + str(i) + "/" + str(len(testfiles)))
tokens = testfiles[i].split("/")[1].split()
testfiles[i] = testfiles[i].replace("/HandStandPushups/",
"/HandstandPushups/")
vid_path = data_path + testfiles[i].split()[0]
vid = imageio.get_reader(vid_path, "ffmpeg")
folder_name = vid_path.split("/")[-1].split(".")[0]
vid_names.append(folder_name)
vid = imageio.get_reader(vid_path, "ffmpeg")
savedir = "../results/images/UCF101/" + prefix + "/" + str(i + 1)
seq_batch = np.zeros(
(1, image_size[0], image_size[1], K + T, c_dim),
dtype="float32")
diff_batch = np.zeros((1, image_size[0], image_size[1], K - 1, 1),
dtype="float32")
for t in range(K + T):
img = vid.get_data(t)[:, :, ::-1]
seq_batch[0, :, :, t] = transform(img)
for t in range(1, K):
prev = inverse_transform(seq_batch[0, :, :, t - 1]) * 255
prev = cv2.cvtColor(prev.astype("uint8"), cv2.COLOR_BGR2GRAY)
next = inverse_transform(seq_batch[0, :, :, t]) * 255
next = cv2.cvtColor(next.astype("uint8"), cv2.COLOR_BGR2GRAY)
diff = next.astype("float32") - prev.astype("float32")
diff_batch[0, :, :, t - 1] = diff[:, :, None] / 255.
true_data = seq_batch[:, :, :, K:, :].copy()
pred_data = np.zeros(true_data.shape, dtype="float32")
xt = seq_batch[:, :, :, K - 1]
pred_data[0] = sess.run(model.G,
feed_dict={
model.diff_in: diff_batch,
model.xt: xt
})
if not os.path.exists(savedir):
os.makedirs(savedir)
cpsnr = np.zeros((K + T, ))
cssim = np.zeros((K + T, ))
pred_data = np.concatenate((seq_batch[:, :, :, :K], pred_data),
axis=3)
true_data = np.concatenate((seq_batch[:, :, :, :K], true_data),
axis=3)
for t in range(K + T):
pred = (inverse_transform(pred_data[0, :, :, t]) *
255).astype("uint8")
target = (inverse_transform(true_data[0, :, :, t]) *
255).astype("uint8")
cpsnr[t] = measure.compare_psnr(pred, target)
cssim[t] = ssim.compute_ssim(Image.fromarray(target),
Image.fromarray(pred))
pred = draw_frame(pred, t < K)
target = draw_frame(target, t < K)
cv2.imwrite(savedir + "/pred_" + "{0:04d}".format(t) + ".png",
pred)
cv2.imwrite(savedir + "/gt_" + "{0:04d}".format(t) + ".png",
target)
cmd1 = "rm " + savedir + "/pred.gif"
cmd2 = ("ffmpeg -f image2 -framerate 3 -i " + savedir +
"/pred_%04d.png " + savedir + "/pred.gif")
cmd3 = "rm " + savedir + "/pred*.png"
# Comment out "system(cmd3)" if you want to keep the output images
# Otherwise only the gifs will be kept
system(cmd1)
system(cmd2)
system(cmd3)
cmd1 = "rm " + savedir + "/gt.gif"
cmd2 = ("ffmpeg -f image2 -framerate 3 -i " + savedir +
"/gt_%04d.png " + savedir + "/gt.gif")
cmd3 = "rm " + savedir + "/gt*.png"
# Comment out "system(cmd3)" if you want to keep the output images
# Otherwise only the gifs will be kept
system(cmd1)
system(cmd2)
system(cmd3)
psnr_err = np.concatenate((psnr_err, cpsnr[None, K:]), axis=0)
ssim_err = np.concatenate((ssim_err, cssim[None, K:]), axis=0)
np.savez(save_path, psnr=psnr_err, ssim=ssim_err)
print("Results saved to " + save_path)
print("Done.")
def main(lr, batch_size, alpha, beta, image_size_h, image_size_w, K,
T, num_iter, gpu):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu[0])
data_path = "/data1/ikusyou/vimeo_septuplet/sequences/"
f = open(data_path+"sep_trainlist.txt","r")
trainfiles = [l[:-1] for l in f.readlines()]
f.close()
# new part for validation files, @@@ please fill the path
val_path = ""
f = open(val_path+"", "r")
val_files = [l[:-1] for l in f.readlines()]
f.close()
val_num = len(val_files)
val_seq = np.zeros((val_num, 240, 416, K+T, 3), dtype="float32")
val_diff = np.zeros((val_num, 240, 416, K-1, 3), dtype="float32")
# At the very beginning, we will read the whole validation set to the memory
for idx, f_name in enumerate(val_files):
output = load_vimeo_data(f_name, val_path, [240, 416], K, T)
val_seq[idx] = output[0]
val_diff[idx] = output[1]
# Now we finish reading the val data
margin = 0.3
updateD = False
updateG = True
iters = 0
prefix = ("VIMEO_MCNET_V1.1.1"
+ "_image_size_h="+str(image_size_h)
+ "_image_size_w="+str(image_size_w)
+ "_K="+str(K)
+ "_T="+str(T)
+ "_batch_size="+str(batch_size)
+ "_lr="+str(lr))
print("\n"+prefix+"\n")
checkpoint_dir = "../models/"+prefix+"/"
samples_dir = "../samples/"+prefix+"/"
summary_dir = "../logs/"+prefix+"/"
if not exists(checkpoint_dir):
makedirs(checkpoint_dir)
if not exists(samples_dir):
makedirs(samples_dir)
if not exists(summary_dir):
makedirs(summary_dir)
with tf.device("/gpu:%d"%gpu[0]):
model = MCNET(image_size=[image_size_h,image_size_w], c_dim=3,
K=K, batch_size=batch_size, T=T,
checkpoint_dir=checkpoint_dir)
# d_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
# model.d_loss, var_list=model.d_vars
# )
g_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
alpha*model.L_img, var_list=model.g_vars
)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False,
gpu_options=gpu_options)) as sess:
tf.global_variables_initializer().run()
if model.load(sess, checkpoint_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
g_sum = tf.summary.merge([model.L_p_sum,
model.L_gdl_sum, model.loss_sum])
# d_sum = tf.summary.merge([model.d_loss_real_sum, model.d_loss_sum,
# model.d_loss_fake_sum])
writer = tf.summary.FileWriter(summary_dir, sess.graph)
counter = iters+1
start_time = time.time()
# IPython.embed()
with Parallel(n_jobs=batch_size) as parallel:
while iters < num_iter:
mini_batches = get_minibatches_idx(len(trainfiles), batch_size, shuffle=True)
for _, batchidx in mini_batches:
if len(batchidx) == batch_size:
seq_batch = np.zeros((batch_size, image_size_h, image_size_w,
K+T, 3), dtype="float32")
diff_batch = np.zeros((batch_size, image_size_h, image_size_w,
K-1, 3), dtype="float32")
t0 = time.time()
Ts = np.repeat(np.array([T]),batch_size,axis=0)
Ks = np.repeat(np.array([K]),batch_size,axis=0)
paths = np.repeat(data_path, batch_size,axis=0)
tfiles = np.array(trainfiles)[batchidx]
shapes = np.repeat(np.array([image_size_h]),batch_size,axis=0)
output = parallel(delayed(load_vimeo_data)(f, p,img_sze, k, t)
for f,p,img_sze,k,t in zip(tfiles,
paths,
shapes,
Ks, Ts))
for i in range(batch_size):
seq_batch[i] = output[i][0]
diff_batch[i] = output[i][1]
_, summary_str = sess.run([g_optim, g_sum],
feed_dict={model.diff_in: diff_batch,
model.xt: seq_batch[:,:,:,K-1],
model.target: seq_batch})
writer.add_summary(summary_str, counter)
errL_img = model.L_img.eval({model.diff_in: diff_batch,
model.xt: seq_batch[:,:,:,K-1],
model.target: seq_batch})
counter += 1
if counter % 50 == 0:
print(
"Iters: [%2d] time: %4.4f, img_loss:%.8f"
% (iters, time.time() - start_time, errL_img)
)
val_loss = []
for idx in range(0, val_num, batch_size)[:-1]: # [:-1] to avoid out if index
diff_batch = val_diff[idx:idx + batch_size]
seq_batch = val_seq[idx:idx + batch_size]
errL_img = model.L_img.eval({model.diff_in: diff_batch,
model.xt: seq_batch[:,:,:,K-1],
model.target: seq_batch})
val_loss.append(np.mean(errL_img))
if np.mod(counter, 200) == 1:
samples = sess.run([model.G],
feed_dict={model.diff_in: diff_batch,
model.xt: seq_batch[:,:,:,K-1],
model.target: seq_batch})[0]
# IPython.embed()
samples = samples[0].swapaxes(0,2).swapaxes(1,2)
# IPython.embed()
sbatch = seq_batch[0,:,:,:].swapaxes(0,2).swapaxes(1,2)
sbatch2 = sbatch.copy()
# IPython.embed()
sbatch2[K:,:,:] = samples
# IPython.embed()
samples = np.concatenate((sbatch2,sbatch), axis=0)
# IPython.embed()
print("Saving sample ...")
save_images(samples, [2, K+T],
samples_dir+"train_%s.png" % (iters))
if np.mod(counter, 10000) == 2:
model.save(sess, checkpoint_dir, counter)
iters += 1
def main(name, lr, batch_size, alpha, beta, image_size, K, T, num_iter, gpu,
nonlinearity, samples_every, gdl, channels, dataset, residual,
planes):
margin = 0.3
updateD = True
updateG = True
iters = 0
namestr = name if len(name) == 0 else "_" + name
prefix = (dataset.replace('/', '-') + namestr + "_channels=" +
str(channels) + "_alpha=" + str(alpha) + "_planes=" +
str(planes) + "_beta=" + str(beta) + "_lr=" + str(lr) +
"_nonlin=" + str(nonlinearity) + "_res=" + str(residual) +
"_gdl=" + str(gdl))
print("\n" + prefix + "\n")
checkpoint_dir = "../models/" + prefix + "/"
samples_dir = "../samples/" + prefix + "/"
summary_dir = "../logs/" + prefix + "/"
normalizer = partial(normalize_data, dataset, channels, K)
train_data, test_data, num_workers = load_dataset(dataset,
T + K,
image_size,
channels,
transforms=[normalizer])
train_loader = DataLoader(train_data,
num_workers=num_workers,
batch_size=batch_size,
shuffle=True,
drop_last=True,
pin_memory=True)
def get_training_batch():
while True:
for sequence in train_loader:
yield sequence
training_batch_generator = get_training_batch()
checkpoint_dir = "../models/" + prefix + "/"
samples_dir = "../samples/" + prefix + "/"
summary_dir = "../logs/" + prefix + "/"
if not exists(checkpoint_dir):
makedirs(checkpoint_dir)
if not exists(samples_dir):
makedirs(samples_dir)
if not exists(summary_dir):
makedirs(summary_dir)
with tf.device("/gpu:%d" % gpu[0]):
model = MCNET(image_size=[image_size, image_size],
c_dim=channels,
K=K,
batch_size=batch_size,
T=T,
checkpoint_dir=checkpoint_dir,
nonlinearity=nonlinearity,
gdl_weight=gdl,
residual=residual,
planes=planes)
d_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
model.d_loss, var_list=model.d_vars)
# facd_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
# model.facd_loss, var_list=model.facd_vars
# )
g_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(
# alpha * model.L_img + beta * model.L_GAN + gamma * model.L_FAC,
alpha * model.L_img + beta * model.L_GAN,
var_list=model.g_vars)
print("GDL: ", model.gdl_weight)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=True,
intra_op_parallelism_threads=3,
inter_op_parallelism_threads=3,
gpu_options=gpu_options))
tf.global_variables_initializer().run()
if model.load(sess, checkpoint_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
g_sum = tf.summary.merge(
[model.L_p_sum, model.L_gdl_sum, model.loss_sum, model.L_GAN_sum])
d_sum = tf.summary.merge(
[model.d_loss_real_sum, model.d_loss_sum, model.d_loss_fake_sum])
writer = tf.summary.FileWriter(summary_dir, sess.graph)
counter = iters + 1
start_time = time.time()
with Parallel(n_jobs=batch_size) as parallel:
while iters < num_iter:
for batch_index in range(100000):
seq_batch, diff_batch = next(training_batch_generator)
# show(seq_batch[0])
# show(diff_batch[0])
# ipdb.set_trace()
seq_batch = seq_batch.numpy()
diff_batch = diff_batch.numpy()
if updateD:
_, summary_str = sess.run(
[d_optim, d_sum],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
writer.add_summary(summary_str, counter)
if updateG:
_, summary_str = sess.run(
[g_optim, g_sum],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
writer.add_summary(summary_str, counter)
errD_fake = model.d_loss_fake.eval({
model.diff_in:
diff_batch,
model.xt:
seq_batch[:, :, :, K - 1],
model.target:
seq_batch
})
errD_real = model.d_loss_real.eval({
model.diff_in:
diff_batch,
model.xt:
seq_batch[:, :, :, K - 1],
model.target:
seq_batch
})
errG = model.L_GAN.eval({
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
errImage = model.L_img.eval({
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})
if errD_fake < margin or errD_real < margin:
updateD = False
if errD_fake > (1. - margin) or errD_real > (1. - margin):
updateG = False
if not updateD and not updateG:
updateD = True
updateG = True
counter += 1
print(
"Iters: [%2d] time: %4.4f, d_loss: %.8f, L_GAN: %.8f, L_img: %.8f"
% (iters, time.time() - start_time, errD_fake + errD_real,
errG, errImage))
if (counter % samples_every) == 0:
samples = sess.run(
[model.G],
feed_dict={
model.diff_in: diff_batch,
model.xt: seq_batch[:, :, :, K - 1],
model.target: seq_batch
})[0]
# ipdb.set_trace()
samples_pad = np.array(samples)
samples_pad.fill(0)
generations = np.concatenate((samples_pad, samples), 3)
generations = generations[0].swapaxes(0, 2).swapaxes(1, 2)
sbatch = seq_batch[0].swapaxes(0, 2).swapaxes(1, 2)
generations = np.concatenate((generations, sbatch), axis=0)
print("Saving sample ...")
# ipdb.set_trace()
save_images(generations[:, :, :, ::-1], [2, T + K],
samples_dir + "train_%s.png" % (iters))
if np.mod(counter, 500) == 2:
model.save(sess, checkpoint_dir, counter)
iters += 1
sess.close()