def testScalarSummaryNameScope(self):
"""Test record_summaries_every_n_global_steps and all_summaries()."""
with ops.Graph().as_default(), self.cached_session() as sess:
global_step = training_util.get_or_create_global_step()
global_step.initializer.run()
with ops.device('/cpu:0'):
step_increment = state_ops.assign_add(global_step, 1)
sess.run(step_increment) # Increment global step from 0 to 1
logdir = tempfile.mkdtemp()
with summary_ops.create_file_writer(logdir, max_queue=0,
name='t2').as_default():
with summary_ops.record_summaries_every_n_global_steps(2):
summary_ops.initialize()
with ops.name_scope('scope'):
summary_op = summary_ops.scalar('my_scalar', 2.0)
# Neither of these should produce a summary because
# global_step is 1 and "1 % 2 != 0"
sess.run(summary_ops.all_summary_ops())
sess.run(summary_op)
events = summary_test_util.events_from_logdir(logdir)
self.assertEqual(len(events), 1)
# Increment global step from 1 to 2 and check that the summary
# is now written
sess.run(step_increment)
sess.run(summary_ops.all_summary_ops())
events = summary_test_util.events_from_logdir(logdir)
self.assertEqual(len(events), 2)
self.assertEqual(events[1].summary.value[0].tag, 'scope/my_scalar')
def testScalarSummaryNameScope(self):
"""Test record_summaries_every_n_global_steps and all_summaries()."""
with ops.Graph().as_default(), self.cached_session() as sess:
global_step = training_util.get_or_create_global_step()
global_step.initializer.run()
with ops.device('/cpu:0'):
step_increment = state_ops.assign_add(global_step, 1)
sess.run(step_increment) # Increment global step from 0 to 1
logdir = tempfile.mkdtemp()
with summary_ops.create_file_writer(logdir, max_queue=0,
name='t2').as_default():
with summary_ops.record_summaries_every_n_global_steps(2):
summary_ops.initialize()
with ops.name_scope('scope'):
summary_op = summary_ops.scalar('my_scalar', 2.0)
# Neither of these should produce a summary because
# global_step is 1 and "1 % 2 != 0"
sess.run(summary_ops.all_summary_ops())
sess.run(summary_op)
events = summary_test_util.events_from_logdir(logdir)
self.assertEqual(len(events), 1)
# Increment global step from 1 to 2 and check that the summary
# is now written
sess.run(step_increment)
sess.run(summary_ops.all_summary_ops())
events = summary_test_util.events_from_logdir(logdir)
self.assertEqual(len(events), 2)
self.assertEqual(events[1].summary.value[0].tag,
'scope/my_scalar')
def _host_call_fn(**kwargs):
"""Training host call.
Creates summaries for training metrics.
Args:
**kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
contain key "global_step" with value of current global_step Tensor.
Returns:
List of summary ops to run on the CPU host.
"""
from tensorflow.python.ops import summary_ops_v2 # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
gs = tf.cast(kwargs.pop("global_step")[0], dtype=tf.int64)
for i, summary in enumerate(current_iteration.summaries):
with summary_ops_v2.create_file_writer(
summary.logdir).as_default():
with summary_ops_v2.record_summaries_every_n_global_steps(
n=self.config.save_summary_steps, global_step=gs):
for j, summary_fn in enumerate(summary_fns[i]):
tensor = kwargs["summary_{}_{}".format(i, j)]
summary_fn(tensor, step=gs)
summary.clear_summary_tuples()
return tf.compat.v1.summary.all_v2_summary_ops()
def construct(self, args):
with self.session.graph.as_default():
# Inputs
self.images = tf.placeholder(tf.float32,
[None, MNIST.H, MNIST.W, MNIST.C],
name="images")
self.labels = tf.placeholder(tf.int64, [None], name="labels")
# Computation
hidden = tf.keras.layers.Flatten()(self.images)
# TODO: Add `args.layers` number of hidden layers with size `args.hidden_layer`,
# using activation from `args.activation`, allowing "none", "relu", "tanh", "sigmoid".
# Store the results back to `hidden` variable.
output_layer = tf.keras.layers.Dense(MNIST.LABELS)(hidden)
self.predictions = tf.argmax(output_layer, axis=1)
# Training
loss = tf.keras.losses.sparse_categorical_crossentropy(
self.labels, output_layer, from_logits=True)
global_step = tf.train.create_global_step()
self.training = tf.train.AdamOptimizer().minimize(
loss, global_step=global_step, name="training")
# Summaries
accuracy = tf.math.reduce_mean(
tf.cast(tf.equal(self.labels, self.predictions), tf.float32))
confusion_matrix = tf.reshape(
tf.confusion_matrix(self.labels,
self.predictions,
weights=tf.not_equal(
self.labels, self.predictions),
dtype=tf.float32),
[1, MNIST.LABELS, MNIST.LABELS, 1])
summary_writer = tf_summary.create_file_writer(args.logdir,
flush_millis=10 *
1000)
self.summaries = {}
with summary_writer.as_default(
), tf_summary.record_summaries_every_n_global_steps(100):
self.summaries["train"] = [
tf_summary.scalar("train/loss", loss),
tf_summary.scalar("train/accuracy", accuracy)
]
with summary_writer.as_default(
), tf_summary.always_record_summaries():
for dataset in ["dev", "test"]:
self.summaries[dataset] = [
tf_summary.scalar(dataset + "/accuracy", accuracy),
tf_summary.image(dataset + "/confusion_matrix",
confusion_matrix)
]
with tf.control_dependencies(self.summaries[dataset]):
self.summaries[dataset].append(summary_writer.flush())
# Initialize variables
self.session.run(tf.global_variables_initializer())
with summary_writer.as_default():
tf_summary.initialize(session=self.session,
graph=self.session.graph)
def testRecordEveryNGlobalSteps(self):
step = training_util.get_or_create_global_step()
logdir = tempfile.mkdtemp()
def run_step():
summary_ops.scalar('scalar', i, step=step)
step.assign_add(1)
with summary_ops.create_file_writer(logdir).as_default(
), summary_ops.record_summaries_every_n_global_steps(2, step):
for i in range(10):
run_step()
# And another 10 steps as a graph function.
run_step_fn = function.defun(run_step)
for i in range(10):
run_step_fn()
events = summary_test_util.events_from_logdir(logdir)
self.assertEqual(len(events), 11)
def train():
parser = config_parser()
args = parser.parse_args()
if args.random_seed is not None:
print('Fixing random seed', args.random_seed)
np.random.seed(args.random_seed)
tf.compat.v1.set_random_seed(args.random_seed)
# Load data
if args.dataset_type == 'llff':
images, poses, bds, render_poses, i_test = load_llff_data(
args.datadir,
args.factor,
recenter=True,
bd_factor=.75,
spherify=args.spherify)
hwf = poses[0, :3, -1]
poses = poses[:, :3, :4]
print('Loaded llff', images.shape, render_poses.shape, hwf,
args.datadir)
if not isinstance(i_test, list):
i_test = [i_test]
if args.llffhold > 0:
print('Auto LLFF holdout,', args.llffhold)
i_test = np.arange(images.shape[0])[::args.llffhold]
i_val = i_test
i_train = np.array([
i for i in np.arange(int(images.shape[0]))
if (i not in i_test and i not in i_val)
])
print('DEFINING BOUNDS')
if args.no_ndc:
near = tf.reduce_min(bds) * .9
far = tf.reduce_max(bds) * 1.
else:
near = 0.
far = 1.
print('NEAR FAR', near, far)
elif args.dataset_type == 'blender':
images, poses, render_poses, hwf, i_split = load_blender_data(
args.datadir, args.half_res, args.testskip)
print('Loaded blender', images.shape, render_poses.shape, hwf,
args.datadir)
i_train, i_val, i_test = i_split
near = 2.
far = 6.
if args.white_bkgd:
images = images[..., :3] * images[..., -1:] + (1. -
images[..., -1:])
else:
images = images[..., :3]
elif args.dataset_type == 'deepvoxels':
images, poses, render_poses, hwf, i_split = load_dv_data(
scene=args.shape, basedir=args.datadir, testskip=args.testskip)
print('Loaded deepvoxels', images.shape, render_poses.shape, hwf,
args.datadir)
i_train, i_val, i_test = i_split
hemi_R = np.mean(np.linalg.norm(poses[:, :3, -1], axis=-1))
near = hemi_R - 1.
far = hemi_R + 1.
else:
print('Unknown dataset type', args.dataset_type, 'exiting')
return
# Cast intrinsics to right types
H, W, focal = hwf
H, W = int(H), int(W)
hwf = [H, W, focal]
if args.render_test:
render_poses = np.array(poses[i_test])
# Create log dir and copy the config file
basedir = args.basedir
expname = args.expname
os.makedirs(os.path.join(basedir, expname), exist_ok=True)
f = os.path.join(basedir, expname, 'args.txt')
with open(f, 'w') as file:
for arg in sorted(vars(args)):
attr = getattr(args, arg)
file.write('{} = {}\n'.format(arg, attr))
if args.config is not None:
f = os.path.join(basedir, expname, 'config.txt')
with open(f, 'w') as file:
file.write(open(args.config, 'r').read())
# Create nerf model
render_kwargs_train, render_kwargs_test, start, grad_vars, models = create_nerf(
args)
bds_dict = {
'near': tf.cast(near, tf.float32),
'far': tf.cast(far, tf.float32),
}
render_kwargs_train.update(bds_dict)
render_kwargs_test.update(bds_dict)
# Short circuit if only rendering out from trained model
if args.render_only:
print('RENDER ONLY')
if args.render_test:
# render_test switches to test poses
images = images[i_test]
else:
# Default is smoother render_poses path
images = None
testsavedir = os.path.join(
basedir, expname, 'renderonly_{}_{:06d}'.format(
'test' if args.render_test else 'path', start))
os.makedirs(testsavedir, exist_ok=True)
print('test poses shape', render_poses.shape)
rgbs, _ = render_path(render_poses,
hwf,
args.chunk,
render_kwargs_test,
gt_imgs=images,
savedir=testsavedir,
render_factor=args.render_factor)
print('Done rendering', testsavedir)
imageio.mimwrite(os.path.join(testsavedir, 'video.mp4'),
to8b(rgbs),
fps=30,
quality=8)
return
# Create optimizer
lrate = args.lrate
if args.lrate_decay > 0:
lrate = tf.keras.optimizers.schedules.ExponentialDecay(
lrate, decay_steps=args.lrate_decay * 1000, decay_rate=0.1)
optimizer = tf.keras.optimizers.Adam(lrate)
models['optimizer'] = optimizer
global_step = tf.compat.v1.train.get_or_create_global_step()
global_step.assign(start)
# Prepare raybatch tensor if batching random rays
N_rand = args.N_rand
use_batching = not args.no_batching
if use_batching:
# For random ray batching.
#
# Constructs an array 'rays_rgb' of shape [N*H*W, 3, 3] where axis=1 is
# interpreted as,
# axis=0: ray origin in world space
# axis=1: ray direction in world space
# axis=2: observed RGB color of pixel
print('get rays')
# get_rays_np() returns rays_origin=[H, W, 3], rays_direction=[H, W, 3]
# for each pixel in the image. This stack() adds a new dimension.
rays = [get_rays_np(H, W, focal, p) for p in poses[:, :3, :4]]
rays = np.stack(rays, axis=0) # [N, ro+rd, H, W, 3]
print('done, concats')
# [N, ro+rd+rgb, H, W, 3]
rays_rgb = np.concatenate([rays, images[:, None, ...]], 1)
# [N, H, W, ro+rd+rgb, 3]
rays_rgb = np.transpose(rays_rgb, [0, 2, 3, 1, 4])
rays_rgb = np.stack([rays_rgb[i] for i in i_train],
axis=0) # train images only
# [(N-1)*H*W, ro+rd+rgb, 3]
rays_rgb = np.reshape(rays_rgb, [-1, 3, 3])
rays_rgb = rays_rgb.astype(np.float32)
print('shuffle rays')
np.random.shuffle(rays_rgb)
print('done')
i_batch = 0
N_iters = 1000000
print('Begin')
print('TRAIN views are', i_train)
print('TEST views are', i_test)
print('VAL views are', i_val)
# Summary writers
writer = tf.summary.create_file_writer(
os.path.join(basedir, 'summaries', expname))
writer.set_as_default()
for i in range(start, N_iters):
time0 = time.time()
# Sample random ray batch
if use_batching:
# Random over all images
batch = rays_rgb[i_batch:i_batch + N_rand] # [B, 2+1, 3*?]
batch = tf.transpose(batch, [1, 0, 2])
# batch_rays[i, n, xyz] = ray origin or direction, example_id, 3D position
# target_s[n, rgb] = example_id, observed color.
batch_rays, target_s = batch[:2], batch[2]
i_batch += N_rand
if i_batch >= rays_rgb.shape[0]:
np.random.shuffle(rays_rgb)
i_batch = 0
else:
# Random from one image
img_i = np.random.choice(i_train)
target = images[img_i]
pose = poses[img_i, :3, :4]
if N_rand is not None:
rays_o, rays_d = get_rays(H, W, focal, pose)
if i < args.precrop_iters:
dH = int(H // 2 * args.precrop_frac)
dW = int(W // 2 * args.precrop_frac)
coords = tf.stack(
tf.meshgrid(tf.range(H // 2 - dH, H // 2 + dH),
tf.range(W // 2 - dW, W // 2 + dW),
indexing='ij'), -1)
if i < 10:
print('precrop', dH, dW, coords[0, 0], coords[-1, -1])
else:
coords = tf.stack(
tf.meshgrid(tf.range(H), tf.range(W), indexing='ij'),
-1)
coords = tf.reshape(coords, [-1, 2])
select_inds = np.random.choice(coords.shape[0],
size=[N_rand],
replace=False)
select_inds = tf.gather_nd(coords, select_inds[:, tf.newaxis])
rays_o = tf.gather_nd(rays_o, select_inds)
rays_d = tf.gather_nd(rays_d, select_inds)
batch_rays = tf.stack([rays_o, rays_d], 0)
target_s = tf.gather_nd(target, select_inds)
##### Core optimization loop #####
with tf.GradientTape() as tape:
# Make predictions for color, disparity, accumulated opacity.
rgb, disp, acc, extras = render(H,
W,
focal,
chunk=args.chunk,
rays=batch_rays,
verbose=i < 10,
retraw=True,
**render_kwargs_train)
# Compute MSE loss between predicted and true RGB.
img_loss = img2mse(rgb, target_s)
trans = extras['raw'][..., -1]
loss = img_loss
psnr = mse2psnr(img_loss)
# Add MSE loss for coarse-grained model
if 'rgb0' in extras:
img_loss0 = img2mse(extras['rgb0'], target_s)
loss += img_loss0
psnr0 = mse2psnr(img_loss0)
gradients = tape.gradient(loss, grad_vars)
optimizer.apply_gradients(zip(gradients, grad_vars))
dt = time.time() - time0
##### end #####
# Rest is logging
def save_weights(net, prefix, i):
path = os.path.join(basedir, expname,
'{}_{:06d}.npy'.format(prefix, i))
np.save(path, net.get_weights())
print('saved weights at', path)
if i % args.i_weights == 0:
for k in models:
save_weights(models[k], k, i)
if i % args.i_video == 0 and i > 0:
rgbs, disps = render_path(render_poses, hwf, args.chunk,
render_kwargs_test)
print('Done, saving', rgbs.shape, disps.shape)
moviebase = os.path.join(basedir, expname,
'{}_spiral_{:06d}_'.format(expname, i))
imageio.mimwrite(moviebase + 'rgb.mp4',
to8b(rgbs),
fps=30,
quality=8)
imageio.mimwrite(moviebase + 'disp.mp4',
to8b(disps / np.max(disps)),
fps=30,
quality=8)
if args.use_viewdirs:
render_kwargs_test['c2w_staticcam'] = render_poses[0][:3, :4]
rgbs_still, _ = render_path(render_poses, hwf, args.chunk,
render_kwargs_test)
render_kwargs_test['c2w_staticcam'] = None
imageio.mimwrite(moviebase + 'rgb_still.mp4',
to8b(rgbs_still),
fps=30,
quality=8)
if i % args.i_testset == 0 and i > 0:
testsavedir = os.path.join(basedir, expname,
'testset_{:06d}'.format(i))
os.makedirs(testsavedir, exist_ok=True)
print('test poses shape', poses[i_test].shape)
render_path(poses[i_test],
hwf,
args.chunk,
render_kwargs_test,
gt_imgs=images[i_test],
savedir=testsavedir)
print('Saved test set')
if i % args.i_print == 0 or i < 10:
print(expname, i, psnr.numpy(), loss.numpy(), global_step.numpy())
print('iter time {:.05f}'.format(dt))
with record_summaries_every_n_global_steps(args.i_print):
tf.summary.scalar('loss', loss)
tf.summary.scalar('psnr', psnr)
tf.summary.histogram('tran', trans)
if args.N_importance > 0:
tf.summary.scalar('psnr0', psnr0)
if i % args.i_img == 0:
# Log a rendered validation view to Tensorboard
img_i = np.random.choice(i_val)
target = images[img_i]
pose = poses[img_i, :3, :4]
rgb, disp, acc, extras = render(H,
W,
focal,
chunk=args.chunk,
c2w=pose,
**render_kwargs_test)
psnr = mse2psnr(img2mse(rgb, target))
# Save out the validation image for Tensorboard-free monitoring
testimgdir = os.path.join(basedir, expname, 'tboard_val_imgs')
if i == 0:
os.makedirs(testimgdir, exist_ok=True)
imageio.imwrite(
os.path.join(testimgdir, '{:06d}.png'.format(i)),
to8b(rgb))
with record_summaries_every_n_global_steps(args.i_img):
tf.summary.image('rgb', to8b(rgb)[tf.newaxis])
tf.summary.image('disp', disp[tf.newaxis, ..., tf.newaxis])
tf.summary.image('acc', acc[tf.newaxis, ..., tf.newaxis])
tf.summary.scalar('psnr_holdout', psnr)
tf.summary.image('rgb_holdout', target[tf.newaxis])
if args.N_importance > 0:
with record_summaries_every_n_global_steps(args.i_img):
tf.summary.image('rgb0',
to8b(extras['rgb0'])[tf.newaxis])
tf.summary.image(
'disp0', extras['disp0'][tf.newaxis, ...,
tf.newaxis])
tf.summary.image(
'z_std', extras['z_std'][tf.newaxis, ...,
tf.newaxis])
global_step.assign_add(1)
