def write(self, obj, names=None, frame=0):
if struct.isstruct(obj):
obj = _transform_for_writing(obj)
if names is None:
names = struct.names(obj)
values = struct.flatten(obj)
names = struct.flatten(names)
names = [_slugify_filename(name) for name in names]
self.write_sim_frame(values, names, frame)
else:
name = str(names) if names is not None else 'unnamed'
self.write_sim_frame([obj], [name], frame)
def dataset_handle(shape, dtype, frames=None):
"""
Creates a single virtual TensorFlow dataset (iterator_handle) for the given struct.
The dataset is expected to hold contain all fields required for loading the obj given the current context item condition.
From the dataset, graph input tensors are derived and arranged into a struct of the same shape as obj.
If an integer is passed to frames, a list of such structs is created by unstacking the second-outer-most dimension of the dataset.
:param shape: tensor shape or struct of tensor shapes
:param dtype: data type of struct of data types matching shape
:param frames: Number of frames contained in each example of the dataset. Expects shape (batch_size, frames, ...)
:type frames: int or None
:return: list of struct and placeholder.
1. If frames=None: valid struct corresponding to obj. If frames>1: list thereof
2. placeholder for a TensorFlow dataset iterator handle (dtype=string)
:rtype: tuple
"""
shapes = tuple(struct.flatten(shape, leaf_condition=is_static_shape))
if struct.isstruct(dtype):
dtypes = tuple(struct.flatten(dtype))
assert len(dtypes) == len(shapes)
else:
dtypes = [dtype] * len(shapes)
if frames is not None:
shapes = tuple(
[shape[0:1] + (frames, ) + shape[1:] for shape in shapes])
# --- TF Dataset handle from string ---
iterator_handle = tf.placeholder(tf.string,
shape=[],
name='dataset_iterator_handle')
iterator = tf.data.Iterator.from_string_handle(iterator_handle,
output_types=dtypes,
output_shapes=shapes)
next_element = iterator.get_next()
# --- Create resulting struct by splitting `next_element`s ---
if frames is None:
next_element_list = list(next_element)
next_struct = struct.map(lambda _: next_element_list.pop(0),
shape,
leaf_condition=is_static_shape)
else:
# --- Remap structures -> to `frames` long list of structs ---
next_struct = []
for frame_idx in range(frames):
next_element_list = list(next_element)
frame_struct = struct.map(
lambda _: next_element_list.pop(0)[:, frame_idx, ...],
shape,
leaf_condition=is_static_shape)
next_struct.append(frame_struct)
return next_struct, iterator_handle
def gradients(y, xs, grad_y=None):
"""
Compute the analytic gradients using TensorFlow's automatic differentiation.
:param y: tensor or struct of tensors. The contributions of all tensors in `y` are added up.
:param xs: struct of input tensors
:return: struct compatible with `xs` holding dy/dx
"""
ys = struct.flatten(y)
if grad_y is not None:
grad_y = struct.flatten(grad_y)
for i in range(len(grad_y)):
grad_y[i] = math.cast(grad_y[i], math.dtype(ys[i]))
xs_ = struct.flatten(xs)
grad = tf.gradients(ys, xs_, grad_ys=grad_y)
return struct.unflatten(grad, xs)
def l1_loss(tensor: Tensor, batch_norm=True, reduce_batches=True):
"""
get L1 loss
Args:
tensor: Tensor:
batch_norm: (Default value = True)
reduce_batches: (Default value = True)
Returns:
"""
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(
l1_loss(tensor, batch_norm, reduce_batches)
for tensor in all_tensors)
if reduce_batches:
total_loss = math.sum_(math.abs(tensor))
else:
total_loss = math.sum_(math.abs(tensor),
dim=list(range(1, len(tensor.shape))))
if batch_norm and reduce_batches:
batch_size = tensor.shape.sizes[0]
return math.divide_no_nan(total_loss, math.to_float(batch_size))
else:
return total_loss
def test_paths(self):
obj = {
'Vels': [CenteredGrid(numpy.zeros([1, 4, 1]), box[0:1], name='v')]
}
with struct.unsafe():
names = struct.flatten(
struct.map(lambda attr: attr.path(), obj, trace=True))
self.assertEqual(names[0], 'Vels.0.data')
def test_names(self):
for obj in generate_test_structs():
with struct.unsafe():
names = struct.flatten(
struct.map(lambda attr: attr.name, obj, trace=True))
self.assertGreater(len(names), 0)
for name in names:
self.assertIsInstance(name, str)
def build_graph_input(obj, input_type='placeholder', frames=None):
"""
Create placeholders for tensors in the supplied state.
:param obj: struct or StateProxy
:param input_type: 'placeholder' or 'dataset_handle'
:param frames: Number of input frames. If not None, returns a list of input structs.
:return:
1. Valid state containing or derived from created placeholders or dataset handle
2. dict mapping from placeholders to their default names (using struct.names)
"""
if isinstance(obj, StateProxy):
obj = obj.state
assert struct.isstruct(obj)
# --- Shapes and names ---
writable_obj = _transform_for_writing(obj)
shape = _writing_staticshape(obj)
names = struct.names(writable_obj)
if input_type == 'placeholder':
if frames is not None: raise NotImplementedError()
with _unsafe():
placeholders = placeholder(shape)
graph_in = struct.map(
lambda x: x,
placeholders) # validates fields, splits staggered tensors
return graph_in, {placeholders: names}
elif input_type == 'dataset_handle':
with _unsafe():
dtypes = struct.dtype(writable_obj)
dataset_nodes, iterator_handle = dataset_handle(shape,
dtypes,
frames=frames)
graph_in = struct.map(
lambda x: x,
dataset_nodes) # validates fields, splits staggered tensors
shapes = struct.flatten(struct.staticshape(dataset_nodes),
leaf_condition=is_static_shape)
dtypes = struct.flatten(struct.dtype(dataset_nodes))
return graph_in, {
'names': struct.flatten(names),
'iterator_handle': iterator_handle,
'shapes': shapes,
'dtypes': dtypes,
'frames': frames
}
else:
raise ValueError(input_type)
def test_flatten(self):
for obj in generate_test_structs():
flat = struct.flatten(obj)
self.assertIsInstance(flat, list)
self.assertGreater(len(flat), 0)
for item in flat:
self.assertTrue(not struct.isstruct(item),
'The result of flatten(%s) is not flat.' % obj)
def l_n_loss(tensor, n, batch_norm=True):
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(l_n_loss(tensor, n, batch_norm) for tensor in all_tensors)
total_loss = math.sum(tensor**n) / n
if batch_norm:
batch_size = math.shape(tensor)[0]
return math.div(total_loss, math.to_float(batch_size))
else:
return total_loss
def read(self, obj, frame=0):
if struct.isstruct(obj):
obj = _transform_for_writing(obj)
names = struct.flatten(obj)
if not np.all([isinstance(n, six.string_types) for n in names]):
names = struct.names(obj)
data = struct.map(lambda name: self.read_array(self._filename(name), frame), names)
return data
else:
return self.read_array('unnamed', frame)
def get_field(self, name):
data = self.app.get_field(name)
tensors = struct.flatten(data)
minimum = min([numpy.min(tensor) for tensor in tensors])
maximum = max([numpy.max(tensor) for tensor in tensors])
if name in self.field_minmax:
midpoint = (self.field_minmax[name][1] + self.field_minmax[name][0])/2
minimum = min(minimum, self.field_minmax[name][0]*self.minmax_decay + midpoint * (1-self.minmax_decay))
maximum = max(maximum, self.field_minmax[name][1]*self.minmax_decay + midpoint * (1-self.minmax_decay))
self.field_minmax[name] = (minimum, maximum)
return data
def step(self, field, dt=1.0, velocity=None):
if not isinstance(velocity, Field):
velocity = velocity.velocity
advected = advect.advect(field, velocity,
dt=dt).copied_with(age=field.age + dt)
if self.conserve and isinstance(
field,
(CenteredGrid, StaggeredGrid)) and np.all(
~np.char.equal(struct.flatten(field.extrapolation), 'constant')
): # If field has zero extrapolation, it cannot be conserved
advected = advected.normalized(field)
return advected
def l1_loss(tensor, batch_norm=True, reduce_batches=True):
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(l1_loss(tensor, batch_norm, reduce_batches) for tensor in all_tensors)
if reduce_batches:
total_loss = math.sum(math.abs(tensor))
else:
total_loss = math.sum(math.abs(tensor), axis=list(range(1, len(tensor.shape))))
if batch_norm and reduce_batches:
batch_size = math.shape(tensor)[0]
return math.div(total_loss, math.to_float(batch_size))
else:
return total_loss
def __init__(self, dataset, fields):
self._dataset = dataset
self._index = 0
self._streams = []
self._fields = fields
self.streams = struct.flatten(fields)
for stream in self.streams:
if isinstance(stream, DataStream):
self._streams.append(stream)
else:
self._streams.append(SourceStream(stream))
self._cache = _BatchCache()
self.indexcache = None
self._dataset_changed()
def __init__(self, dataset, fields):
self._dataset = dataset
self._index = 0
self._streams = []
self._fields = fields
self.streams = struct.flatten(fields)
for stream in self.streams:
if isinstance(stream, DataStream):
self._streams.append(stream)
else:
self._streams.append(SourceStream(stream))
self._cache = _BatchCache()
self.indexcache = None
self.callback = self._dataset_changed # Permanent reference so it won't be garbage collected
self._dataset.on_change(self.callback)
self._dataset_changed(self._dataset)
def frequency_loss(tensor, frequency_falloff=100, reduce_batches=True):
"""
Instead of minimizing each entry of the tensor, minimize the frequencies of the tensor, emphasizing lower frequencies over higher ones.
:param reduce_batches: whether to reduce the batch dimension of the loss by adding the losses along the first dimension
:param tensor: typically actual - target
:param frequency_falloff: large values put more emphasis on lower frequencies, 1.0 weights all frequencies equally.
:return: scalar loss value
"""
if struct.isstruct(tensor):
all_tensors = struct.flatten(tensor)
return sum(
frequency_loss(tensor, frequency_falloff, reduce_batches)
for tensor in all_tensors)
diff_fft = abs_square(math.fft(tensor))
k = fftfreq(tensor.shape[1:-1], mode='absolute')
weights = math.exp(-0.5 * k**2 * frequency_falloff**2)
return l1_loss(diff_fft * weights, reduce_batches=reduce_batches)
def __init__(self, dataset, fields):
self._dataset = dataset
self._index = 0
self._streams = []
self._fields = fields
self.streams = tuple(
filter(
lambda x: isinstance(x, DataStream) or isinstance(
x, six.string_types), struct.flatten(fields)))
self.stream_mask = struct.map(lambda x: x in self.streams,
self._fields,
content_type='stream_mask')
for stream in self.streams:
if isinstance(stream, DataStream):
self._streams.append(stream)
elif isinstance(stream, six.string_types):
self._streams.append(SourceStream(stream))
else:
assert False
self._cache = _BatchCache()
self.indexcache = None
self._dataset_changed()
def run(self,
fetches,
feed_dict=None,
summary_key=None,
time=None,
merged_summary=None,
item_condition=struct.ALL_ITEMS):
if isinstance(fetches, np.ndarray):
return fetches
if fetches is None:
return None
tensor_feed_dict = None
if feed_dict is not None:
tensor_feed_dict = {}
for (key, value) in feed_dict.items():
pairs = struct.zip([key, value],
item_condition=item_condition,
zip_parents_if_incompatible=True)
def add_to_dict(key_tensor, value_tensor):
if isplaceholder(key_tensor):
tensor_feed_dict[key_tensor] = value_tensor
return None
struct.map(add_to_dict,
pairs,
item_condition=item_condition,
content_type=struct.INVALID)
tensor_fetches = struct.flatten(fetches, item_condition=item_condition)
if isinstance(fetches, (tuple, list)):
def is_fetch(x):
return istensor(x) or _identity_in(x, fetches)
else:
def is_fetch(x):
return istensor(x) or x is fetches
tensor_fetches = tuple(filter(is_fetch, tensor_fetches))
# Handle tracing
trace = _trace_stack.get_default(raise_error=False)
if trace:
options = trace.timeliner.options
run_metadata = trace.timeliner.run_metadata
else:
options = None
run_metadata = None
# Summary
if summary_key is not None and merged_summary is not None:
tensor_fetches = (merged_summary, ) + tensor_fetches
result_fetches = self._session.run(tensor_fetches, tensor_feed_dict,
options, run_metadata)
result_dict = {
fetch: result
for fetch, result in zip(tensor_fetches, result_fetches)
}
if summary_key:
summary_buffer = result_fetches[0]
result_fetches = result_fetches[1:]
if summary_key in self.summary_writers:
summary_writer = self.summary_writers[summary_key]
else:
summary_writer = tf.summary.FileWriter(
os.path.join(self.summary_directory, str(summary_key)),
self.graph)
self.summary_writers[summary_key] = summary_writer
summary_writer.add_summary(summary_buffer, time)
summary_writer.flush()
if trace:
trace.timeliner.add_run()
def replace_tensor_with_value(fetch):
try:
if fetch in result_dict:
return result_dict[fetch]
else:
return fetch
except TypeError: # not hashable
return fetch
result = struct.map(replace_tensor_with_value,
fetches,
item_condition=item_condition)
return result
