def __init__(self, model_path):
#
tfutil.init_tf()
#
print('[MD] Load insight face model...')
with tf.io.gfile.GFile(model_path, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name="")
self.image_input = tf.get_default_graph().get_tensor_by_name('data:0')
embedding = tf.get_default_graph().get_tensor_by_name('output:0')
embedding_norm = tf.norm(embedding, axis=1, keepdims=True)
self.embedding = tf.div(embedding,
embedding_norm,
name='norm_embedding')
self.target_emb = tf.placeholder(tf.float32,
shape=[None, 512],
name='target_emb_input')
self.cos_loss = tf.reduce_sum(
tf.multiply(self.embedding, self.target_emb))
self.l2_loss = tf.norm(self.embedding - self.target_emb)
self.grads_op = tf.gradients(self.l2_loss, self.image_input)
# self.fdict = {keep_prob:1.0, is_train:False}
self.fdict = {}
self.sess = tf.get_default_session()
def infer_stack(args):
tmp = Image.open(args.stack)
h,w = np.shape(tmp)
N = tmp.n_frames
imgs = np.zeros((N, 3, h, w))
for i in range(N):
tmp.seek(i)
imgs[i, 0, :, :] = np.array(tmp)
imgs[i, 1, :, :] = np.array(tmp)
imgs[i, 2, :, :] = np.array(tmp)
imgs = imgs.astype("float32")
imgs = imgs / 255.0 - 0.5
tfutil.init_tf(tf_config)
net = util.load_snapshot(args.network)
res = np.empty((N, h, w), dtype="uint16")
for i in range(N):
res[i,:,:] = util.infer_image_pp(net, imgs[i,:,:,:])
#tmp = Image.fromarray(res[0,:,:,:].transpose([1,2,0]).astype("uint8"))
tmp = Image.fromarray(res[0,:,:])
tmp.save(args.out, format="tiff",
append_images=[Image.fromarray(res[i,:,:]) for i in range(1, res.shape[0])],
save_all=True)
def init_tf(random_seed=1234):
"""Initialize TF."""
print('Initializing TensorFlow...\n')
np.random.seed(random_seed)
tfutil.init_tf({
'graph_options.place_pruned_graph': True,
'gpu_options.allow_growth': True
})
def infer_image(network_snapshot: str, image: str, out_image: str):
tfutil.init_tf(config.tf_config)
net = util.load_snapshot(network_snapshot)
im = PIL.Image.open(image).convert('RGB')
arr = np.array(im, dtype=np.float32)
reshaped = arr.transpose([2, 0, 1]) / 255.0 - 0.5
pred255 = util.infer_image(net, reshaped)
t = pred255.transpose([1, 2, 0]) # [RGB, H, W] -> [H, W, RGB]
PIL.Image.fromarray(t, 'RGB').save(os.path.join(out_image))
print('Inferred image saved in', out_image)
def validate(submit_config: dnnlib.SubmitConfig, noise: dict, dataset: dict, network_snapshot: str):
noise_augmenter = dnnlib.util.call_func_by_name(**noise)
validation_set = ValidationSet(submit_config)
validation_set.load(**dataset)
ctx = dnnlib.RunContext(submit_config, config)
tfutil.init_tf(config.tf_config)
with tf.device("/gpu:0"):
net = util.load_snapshot(network_snapshot)
validation_set.evaluate(net, 0, noise_augmenter.add_validation_noise_np)
ctx.close()
def validate(submit_config: submit.SubmitConfig, tf_config: dict, noise: dict,
dataset: dict, network_snapshot: str):
noise_augmenter = noise.func(**noise.func_kwargs)
validation_set = ValidationSet(submit_config)
validation_set.load(**dataset)
ctx = RunContext(submit_config)
tfutil.init_tf(tf_config)
with tf.device("/gpu:0"):
net = util.load_snapshot(network_snapshot)
validation_set.evaluate(net, 0,
noise_augmenter.add_validation_noise_np)
ctx.close()
def validate(submit_config: dnnlib.SubmitConfig, noise: dict, dataset: dict,
network_snapshot: str):
noise_augmenter = dnnlib.util.call_func_by_name(**noise)
validation_set = ValidationSet(submit_config)
validation_set.load(**dataset)
# Create a run context (hides low level details, exposes simple API to manage the run)
ctx = dnnlib.RunContext(submit_config, config)
# Initialize TensorFlow graph and session using good default settings
tfutil.init_tf(config.tf_config)
# Construct the network using the Network helper class and a function defined in config.net_config
with tf.device("/gpu:0"):
net = load_snapshot(network_snapshot)
validation_set.evaluate(net, 0,
noise_augmenter.add_validation_noise_np)
ctx.close()
def train(submit_config: dnnlib.SubmitConfig, iteration_count: int,
eval_interval: int, minibatch_size: int, learning_rate: float,
ramp_down_perc: float, noise: dict, validation_config: dict,
train_tfrecords: str, noise2noise: bool):
noise_augmenter = dnnlib.util.call_func_by_name(**noise)
validation_set = ValidationSet(submit_config)
validation_set.load(**validation_config)
# Create a run context (hides low level details, exposes simple API to manage the run)
# noinspection PyTypeChecker
ctx = dnnlib.RunContext(submit_config, config)
# Initialize TensorFlow graph and session using good default settings
tfutil.init_tf(config.tf_config)
dataset_iter = create_dataset(train_tfrecords, minibatch_size,
noise_augmenter.add_train_noise_tf)
# Construct the network using the Network helper class and a function defined in config.net_config
with tf.device("/gpu:0"):
net = tflib.Network(**config.net_config)
# Optionally print layer information
net.print_layers()
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device("/cpu:0"):
lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
noisy_input, noisy_target, clean_target = dataset_iter.get_next()
noisy_input_split = tf.split(noisy_input, submit_config.num_gpus)
noisy_target_split = tf.split(noisy_target, submit_config.num_gpus)
clean_target_split = tf.split(clean_target, submit_config.num_gpus)
# Define the loss function using the Optimizer helper class, this will take care of multi GPU
opt = tflib.Optimizer(learning_rate=lrate_in, **config.optimizer_config)
for gpu in range(submit_config.num_gpus):
with tf.device("/gpu:%d" % gpu):
net_gpu = net if gpu == 0 else net.clone()
denoised = net_gpu.get_output_for(noisy_input_split[gpu])
if noise2noise:
meansq_error = tf.reduce_mean(
tf.square(noisy_target_split[gpu] - denoised))
else:
meansq_error = tf.reduce_mean(
tf.square(clean_target_split[gpu] - denoised))
# Create an autosummary that will average over all GPUs
with tf.control_dependencies([autosummary("Loss", meansq_error)]):
opt.register_gradients(meansq_error, net_gpu.trainables)
train_step = opt.apply_updates()
# Create a log file for Tensorboard
summary_log = tf.summary.FileWriter(submit_config.run_dir)
summary_log.add_graph(tf.get_default_graph())
print('Training...')
time_maintenance = ctx.get_time_since_last_update()
ctx.update(loss='run %d' % submit_config.run_id,
cur_epoch=0,
max_epoch=iteration_count)
# ***********************************
# The actual training loop
for i in range(iteration_count):
# Whether to stop the training or not should be asked from the context
if ctx.should_stop():
break
# Dump training status
if i % eval_interval == 0:
time_train = ctx.get_time_since_last_update()
time_total = ctx.get_time_since_start()
# Evaluate 'x' to draw a batch of inputs
[source_mb, target_mb] = tfutil.run([noisy_input, clean_target])
denoised = net.run(source_mb)
save_image(submit_config, denoised[0],
"img_{0}_y_pred.png".format(i))
save_image(submit_config, target_mb[0], "img_{0}_y.png".format(i))
save_image(submit_config, source_mb[0],
"img_{0}_x_aug.png".format(i))
validation_set.evaluate(net, i,
noise_augmenter.add_validation_noise_np)
print(
'iter %-10d time %-12s eta %-12s sec/eval %-7.1f sec/iter %-7.2f maintenance %-6.1f'
% (autosummary('Timing/iter', i),
dnnlib.util.format_time(
autosummary('Timing/total_sec', time_total)),
dnnlib.util.format_time(
autosummary('Timing/total_sec',
(time_train / eval_interval) *
(iteration_count - i))),
autosummary('Timing/sec_per_eval', time_train),
autosummary('Timing/sec_per_iter',
time_train / eval_interval),
autosummary('Timing/maintenance_sec', time_maintenance)))
dnnlib.tflib.autosummary.save_summaries(summary_log, i)
ctx.update(loss='run %d' % submit_config.run_id,
cur_epoch=i,
max_epoch=iteration_count)
time_maintenance = ctx.get_last_update_interval() - time_train
# Training epoch
lrate = compute_ramped_down_lrate(i, iteration_count, ramp_down_perc,
learning_rate)
tfutil.run([train_step], {lrate_in: lrate})
# End of training
print("Elapsed time: {0}".format(
util.format_time(ctx.get_time_since_start())))
save_snapshot(submit_config, net, 'final')
# Summary log and context should be closed at the end
summary_log.close()
ctx.close()
def train(
submit_config: submit.SubmitConfig,
iteration_count: int,
eval_interval: int,
minibatch_size: int,
learning_rate: float,
ramp_down_perc: float,
noise: dict,
tf_config: dict,
net_config: dict,
optimizer_config: dict,
validation_config: dict,
train_tfrecords: str):
# **dict as argument means: take all additional named arguments to this function
# and insert them into this parameter as dictionary entries.
noise_augmenter = noise.func(**noise.func_kwargs)
validation_set = ValidationSet(submit_config)
# Load all images for validation as numpy arrays to the images attribute of the validation set.
validation_set.load(**validation_config)
# Create a run context (hides low level details, exposes simple API to manage the run)
ctx = run_context.RunContext(submit_config)
# Initialize TensorFlow graph and session using good default settings
tfutil.init_tf(tf_config)
# Creates the data set from the specified path to a generated tfrecords file containing all training images.
# Data set will be split into minibatches of the given size and augment the noise with given noise function.
# Use the dataset_tool_tf to create this tfrecords file.
dataset_iter = create_dataset(train_tfrecords, minibatch_size, noise_augmenter.add_train_noise_tf)
# Construct the network using the Network helper class and a function defined in config.net_config
with tf.device("/gpu:0"):
net = Network(**net_config)
# Optionally print layer information
net.print_layers()
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device("/cpu:0"):
# Placeholder for the learning rate. This will get ramped down dynamically.
lrate_in = tf.placeholder(tf.float32, name='lrate_in', shape=[])
# Defines the expression(s) that creates the network input.
noisy_input, noisy_target, clean_target = dataset_iter.get_next()
noisy_input_split = tf.split(noisy_input, submit_config.num_gpus)
noisy_target_split = tf.split(noisy_target, submit_config.num_gpus) # Split over multiple GPUs
# clean_target_split = tf.split(clean_target, submit_config.num_gpus)
# --------------------------------------------------------------------------------------------
# Optimizer initialization and setup:
# Define the loss function using the Optimizer helper class, this will take care of multi GPU
opt = Optimizer(learning_rate=lrate_in, **optimizer_config)
for gpu in range(submit_config.num_gpus):
with tf.device("/gpu:%d" % gpu):
# Create a clone for this network for other gpus to work on.
net_gpu = net if gpu == 0 else net.clone()
# Create the output expression by giving the input expression into the network.
denoised = net_gpu.get_output_for(noisy_input_split[gpu])
# Create the error function as the MSE between the target tensor and the denoised network output.
meansq_error = tf.reduce_mean(tf.square(noisy_target_split[gpu] - denoised))
# Create an autosummary that will average over all GPUs
with tf.control_dependencies([autosummary("Loss", meansq_error)]):
opt.register_gradients(meansq_error, net_gpu.trainables)
train_step = opt.apply_updates() # Defines the update function of the optimizer.
# Create a log file for Tensorboard
summary_log = tf._api.v1.summary.FileWriter(submit_config.results_dir)
summary_log.add_graph(tf.get_default_graph())
# --------------------------------------------------------------------------------------------
# Training and some milestone evaluation starts:
print('Training...')
time_maintenance = ctx.get_time_since_last_update()
ctx.update() # TODO: why parameterized in reference?
# The actual training loop
for i in range(iteration_count):
# Whether to stop the training or not should be asked from the context
if ctx.should_stop():
break
# Dump training status
if i % eval_interval == 0:
time_train = ctx.get_time_since_last_update()
time_total = ctx.get_time_since_start()
# Evaluate 'x' to draw one minbatch of inputs. Executes the operations defined in the dataset iterator.
# Evals the noisy input and clean target minibatch Tensor ops to numpy array of the minibatch.
[source_mb, target_mb] = tfutil.run([noisy_input, clean_target])
# Runs the noisy images through the network without training it. It is just for observing/evaluating.
# net.run expects numpy arrays to run through this network.
denoised = net.run(source_mb)
# array shape: [minibatch_size, channel_size, height, width]
util.save_image(submit_config, denoised[0], "img_{0}_y_pred.png".format(i))
util.save_image(submit_config, target_mb[0], "img_{0}_y.png".format(i))
util.save_image(submit_config, source_mb[0], "img_{0}_x_aug.png".format(i))
validation_set.evaluate(net, i, noise_augmenter.add_validation_noise_np)
print('iter %-10d time %-12s sec/eval %-7.1f sec/iter %-7.2f maintenance %-6.1f' % (
autosummary('Timing/iter', i),
dnnlib.util.format_time(autosummary('Timing/total_sec', time_total)),
autosummary('Timing/sec_per_eval', time_train),
autosummary('Timing/sec_per_iter', time_train / eval_interval),
autosummary('Timing/maintenance_sec', time_maintenance)))
dnnlib.tflib.autosummary.save_summaries(summary_log, i)
ctx.update()
time_maintenance = ctx.get_last_update_interval() - time_train
lrate = compute_ramped_down_lrate(i, iteration_count, ramp_down_perc, learning_rate)
# Apply the lrate value to the lrate_in placeholder for the optimizer.
tfutil.run([train_step], {lrate_in: lrate}) # Run the training update through the network in our session.
print("Elapsed time: {0}".format(dutil.format_time(ctx.get_time_since_start())))
util.save_snapshot(submit_config, net, 'final')
# Summary log and context should be closed at the end
summary_log.close()
ctx.close()
import moviepy.editor # pip install moviepy
moviepy.editor.VideoClip(make_frame,
duration=duration_sec).write_videofile(
os.path.join(result_subdir, mp4),
fps=mp4_fps,
codec='png',
bitrate=mp4_bitrate)
with open(os.path.join(result_subdir, mp4 + '-keyframes.txt'),
'w') as file:
file.write(str(latents_idx))
## -----------------------------------------------------------------------------------------------------------------
if __name__ == "__main__":
import datetime
import time
print(datetime.datetime.now(), int(time.time()))
np.random.seed(int(time.time()))
tfutil.init_tf()
generate_fake_images(00, num_pngs=3500)
#generate_interpolation_video(12, grid_size=[1,1], random_seed=int(time.time()), mp4_fps=25, duration_sec=300.0)
keyframes = [
75, 125, 300, 375, 450, 550, 600, 700, 775, 800, 850, 900, 925, 975,
1075, 1200, 1350, 1450, 1575, 1700, 1800, 1900, 2050, 2150, 2250, 2300,
2375, 2475, 2800, 2875, 3000, 3050, 3125, 3225, 3300, 3400
]
generate_keyframed_video(00, keyframes, mp4_fps=30)
print('Exiting...')
print(datetime.datetime.now())
def main():
starttime = int(time.time())
parser = argparse.ArgumentParser(
description='Render from StyleGAN2 saved models (pkl files).',
#epilog=_examples,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
'--network_pkl',
help='The pkl file to render from (the model checkpoint).',
default=None,
metavar='MODEL.pkl',
required=True)
parser.add_argument(
'--grid_x',
help=
'Number of images to render horizontally (each frame will have rows of X images, default: 1).',
default=1,
metavar='X',
type=int)
parser.add_argument(
'--grid_y',
help=
'Number of images to render vertically (each frame will have cols of Y images, default: 1).',
default=1,
metavar='Y',
type=int)
parser.add_argument(
'--png_sequence',
help=
'If True, outputs a folder of frames as pngs instead of video. (default: False).',
default=False,
type=bool)
parser.add_argument(
'--image_shrink',
help=
'Render in 1/[image_shrink] resolution (fast, useful for quick previews)',
default=1,
type=int)
parser.add_argument(
'--image_zoom',
help=
'Zoom on the output image (seems like just more video pixels, but no true upscaling)',
default=1,
type=float)
parser.add_argument('--duration_sec',
help='Length of video to render in seconds.',
default=30.0,
type=float)
parser.add_argument('--mp4_fps',
help='Frames per second for video rendering',
default=30,
type=float)
parser.add_argument(
'--smoothing_sec',
help=
'Gaussian kernel size in seconds to blend video frames (higher value = less change, lower value = more erratic, default: 1.0)',
default=1.0,
type=float)
parser.add_argument(
'--truncation_psi',
help=
'Truncation parameter (1 = normal, lower values overfit to look more like originals, higher values underfit to be more abstract, recommendation: 0.5-2)',
default=1,
type=float)
parser.add_argument('--randomize_noise',
help='If True, adds noise to vary rendered images.',
default=False,
type=bool)
parser.add_argument(
'--filename',
help='Filename for rendering output, defaults to pkl filename',
default=None)
parser.add_argument(
'--mp4_codec',
help='Video codec to use with moviepy (i.e. libx264, libx265, mpeg4)',
default='libx265')
parser.add_argument('--mp4_bitrate',
help='Bitrate to use with moviepy (i.e. 16M)',
default='16M')
parser.add_argument('--random_seed',
help='Seed to initialize the latent generation.',
default=starttime,
type=int)
parser.add_argument(
'--minibatch_size',
help=
'Size of batch rendering (doesn\'t seem to have effects but left in anyway)',
default=8,
type=int)
args = parser.parse_args()
tfutil.init_tf()
generate_interpolation_video(network_pkl=args.network_pkl,
grid_size=[args.grid_x, args.grid_y],
png_sequence=args.png_sequence,
image_shrink=args.image_shrink,
image_zoom=args.image_zoom,
duration_sec=args.duration_sec,
smoothing_sec=args.smoothing_sec,
truncation_psi=args.truncation_psi,
randomize_noise=args.randomize_noise,
filename=args.filename,
mp4_fps=args.mp4_fps,
mp4_codec=args.mp4_codec,
mp4_bitrate=args.mp4_bitrate,
random_seed=args.random_seed,
minibatch_size=args.minibatch_size)
def train(submit_config: dnnlib.SubmitConfig, iteration_count: int,
eval_interval: int, minibatch_size: int, learning_rate: float,
ramp_down_perc: float, noise: dict, validation_config: dict,
train_tfrecords: str, noise2noise: bool):
noise_augmenter = dnnlib.util.call_func_by_name(**noise)
validation_set = ValidationSet(submit_config)
validation_set.load(**validation_config)
# Create a run context (hides low level details, exposes simple API to manage the run)
ctx = dnnlib.RunContext(submit_config, config)
# Initialize TensorFlow graph and session using good default settings
tfutil.init_tf(config.tf_config)
dataset_iter = create_dataset(train_tfrecords, minibatch_size,
noise_augmenter.add_train_noise_tf)
# Construct the network using the Network helper class and a function defined in config.net_config
with tf.device("/gpu:0"):
net = tflib.Network(**config.net_config)
# Optionally print layer information
net.print_layers()
print('Building TensorFlow graph...')
with tf.name_scope('Inputs'), tf.device("/cpu:0"):
lrate_in = tf.compat.v1.placeholder(tf.float32,
name='lrate_in',
shape=[])
#print("DEBUG train:", "dataset iter got called")
noisy_input, noisy_target, clean_target = dataset_iter.get_next()
noisy_input_split = tf.split(noisy_input, submit_config.num_gpus)
noisy_target_split = tf.split(noisy_target, submit_config.num_gpus)
print(len(noisy_input_split), noisy_input_split)
clean_target_split = tf.split(clean_target, submit_config.num_gpus)
# Split [?, 3, 256, 256] across num_gpus over axis 0 (i.e. the batch)
# Define the loss function using the Optimizer helper class, this will take care of multi GPU
opt = tflib.Optimizer(learning_rate=lrate_in, **config.optimizer_config)
radii = np.arange(128).reshape(128, 1) #image size 256, binning = 3
radial_masks = np.apply_along_axis(radial_mask, 1, radii, 128, 128,
np.arange(0, 256), np.arange(0, 256),
20)
print("RN SHAPE!!!!!!!!!!:", radial_masks.shape)
radial_masks = np.expand_dims(radial_masks, 1) # (128, 1, 256, 256)
#radial_masks = np.squeeze(np.stack((radial_masks,) * 3, -1)) # 43, 3, 256, 256
#radial_masks = radial_masks.transpose([0, 3, 1, 2])
radial_masks = radial_masks.astype(np.complex64)
radial_masks = tf.expand_dims(radial_masks, 1)
rn = tf.compat.v1.placeholder_with_default(radial_masks,
[128, None, 1, 256, 256])
rn_split = tf.split(rn, submit_config.num_gpus, axis=1)
freq_nyq = int(np.floor(int(256) / 2.0))
spatial_freq = radii.astype(np.float32) / freq_nyq
spatial_freq = spatial_freq / max(spatial_freq)
for gpu in range(submit_config.num_gpus):
with tf.device("/gpu:%d" % gpu):
net_gpu = net if gpu == 0 else net.clone()
denoised_1 = net_gpu.get_output_for(noisy_input_split[gpu])
denoised_2 = net_gpu.get_output_for(noisy_target_split[gpu])
print(noisy_input_split[gpu].get_shape(),
rn_split[gpu].get_shape())
if noise2noise:
meansq_error = fourier_ring_correlation(
noisy_target_split[gpu], denoised_1, rn_split[gpu],
spatial_freq) - fourier_ring_correlation(
noisy_target_split[gpu] - denoised_2,
noisy_input_split[gpu] - denoised_1, rn_split[gpu],
spatial_freq)
else:
meansq_error = tf.reduce_mean(
tf.square(clean_target_split[gpu] - denoised))
# Create an autosummary that will average over all GPUs
#tf.summary.histogram(name, var)
with tf.control_dependencies([autosummary("Loss", meansq_error)]):
opt.register_gradients(meansq_error, net_gpu.trainables)
train_step = opt.apply_updates()
# Create a log file for Tensorboard
summary_log = tf.compat.v1.summary.FileWriter(submit_config.run_dir)
summary_log.add_graph(tf.compat.v1.get_default_graph())
print('Training...')
time_maintenance = ctx.get_time_since_last_update()
ctx.update(loss='run %d' % submit_config.run_id,
cur_epoch=0,
max_epoch=iteration_count)
# The actual training loop
for i in range(iteration_count):
# Whether to stop the training or not should be asked from the context
if ctx.should_stop():
break
# Dump training status
if i % eval_interval == 0:
time_train = ctx.get_time_since_last_update()
time_total = ctx.get_time_since_start()
print("DEBUG TRAIN!", noisy_input.dtype, noisy_input[0][0].dtype)
# Evaluate 'x' to draw a batch of inputs
[source_mb, target_mb] = tfutil.run([noisy_input, clean_target])
denoised = net.run(source_mb)
save_image(submit_config, denoised[0],
"img_{0}_y_pred.tif".format(i))
save_image(submit_config, target_mb[0], "img_{0}_y.tif".format(i))
save_image(submit_config, source_mb[0],
"img_{0}_x_aug.tif".format(i))
validation_set.evaluate(net, i,
noise_augmenter.add_validation_noise_np)
print(
'iter %-10d time %-12s sec/eval %-7.1f sec/iter %-7.2f maintenance %-6.1f'
% (autosummary('Timing/iter', i),
dnnlib.util.format_time(
autosummary('Timing/total_sec', time_total)),
autosummary('Timing/sec_per_eval', time_train),
autosummary('Timing/sec_per_iter',
time_train / eval_interval),
autosummary('Timing/maintenance_sec', time_maintenance)))
dnnlib.tflib.autosummary.save_summaries(summary_log, i)
ctx.update(loss='run %d' % submit_config.run_id,
cur_epoch=i,
max_epoch=iteration_count)
time_maintenance = ctx.get_last_update_interval() - time_train
save_snapshot(submit_config, net, str(i))
lrate = compute_ramped_down_lrate(i, iteration_count, ramp_down_perc,
learning_rate)
tfutil.run([train_step], {lrate_in: lrate})
print("Elapsed time: {0}".format(
util.format_time(ctx.get_time_since_start())))
save_snapshot(submit_config, net, 'final')
# Summary log and context should be closed at the end
summary_log.close()
ctx.close()
