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()
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_subsampled(
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 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
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()
def copy_vars_from(self, src_net: "Network") -> None:
"""Copy the values of all variables from the given network, including sub-networks."""
names = [name for name in self.vars.keys() if name in src_net.vars]
tfutil.set_vars(
tfutil.run({self.vars[name]: src_net.vars[name]
for name in names}))
def reset_trainables(self) -> None:
"""Re-initialize all trainable variables of this network, including sub-networks."""
tfutil.run([var.initializer for var in self.trainables.values()])
def reset_own_vars(self) -> None:
"""Re-initialize all variables of this network, excluding sub-networks."""
tfutil.run([var.initializer for var in self.own_vars.values()])
def copy_weights(src_net, tgt_net, vars_to_copy):
names = [name for name in tgt_net.trainables.keys() if name in vars_to_copy]
tfutil.set_vars(tfutil.run({tgt_net.vars[name]: src_net.vars[name] for name in names}))
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()