[Occupancy challenge] Pre-processing with the toolkit inside the dataset

[Hediby]

Hi,

I am trying to create a dataset that would yield a tensor directly, instead of a dict. To do so, I need to create objects and use functions from the toolkit (creating bird-eye views, rasterizing maps, etc.). Here is what I do:

def _make_model_inputs(
    timestep_grids: occupancy_flow_grids.TimestepGrids,
    vis_grids: occupancy_flow_grids.VisGrids,
) -> tf.Tensor:
    """Concatenates all occupancy grids over past, current to a single tensor."""
    model_inputs = tf.concat(
      [
          vis_grids.roadgraph,
          timestep_grids.vehicles.past_occupancy,
          timestep_grids.vehicles.current_occupancy,
          tf.clip_by_value(
              timestep_grids.pedestrians.past_occupancy +
              timestep_grids.cyclists.past_occupancy, 0, 1),
          tf.clip_by_value(
              timestep_grids.pedestrians.current_occupancy +
              timestep_grids.cyclists.current_occupancy, 0, 1),
      ],
      axis=-1,
    )
    return model_inputs

def make_model_inputs(inputs):
    timestep_grids = occupancy_flow_grids.create_ground_truth_timestep_grids(
      inputs, config)
    true_waypoints = occupancy_flow_grids.create_ground_truth_waypoint_grids(
      timestep_grids, config)
    vis_grids = occupancy_flow_grids.create_ground_truth_vis_grids(
      inputs, timestep_grids, config)
    return _make_model_inputs(timestep_grids, vis_grids)

filenames = tf.io.matching_files(TRAIN_FILES)
dataset = tf.data.TFRecordDataset(filenames, compression_type='')
dataset = dataset.map(occupancy_flow_data.parse_tf_example)
dataset = dataset.batch(5)
dataset = dataset.map(make_model_inputs)
it = iter(dataset)
inputs = next(it)

This gives me the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-10-0dff10c5b3ba> in <module>
     33 dataset = dataset.map(occupancy_flow_data.parse_tf_example)
     34 dataset = dataset.batch(5)
---> 35 dataset = dataset.map(make_model_inputs)
     36 it = iter(dataset)
     37 inputs = next(it)

~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in map(self, map_func, num_parallel_calls, deterministic)
   1859         warnings.warn("The `deterministic` argument has no effect unless the "
   1860                       "`num_parallel_calls` argument is specified.")
-> 1861       return MapDataset(self, map_func, preserve_cardinality=True)
   1862     else:
   1863       return ParallelMapDataset(

~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, input_dataset, map_func, use_inter_op_parallelism, preserve_cardinality, use_legacy_function)
   4983         self._transformation_name(),
   4984         dataset=input_dataset,
-> 4985         use_legacy_function=use_legacy_function)
   4986     variant_tensor = gen_dataset_ops.map_dataset(
   4987         input_dataset._variant_tensor,  # pylint: disable=protected-access

~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in __init__(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)
   4216         fn_factory = trace_tf_function(defun_kwargs)
   4217 
-> 4218     self._function = fn_factory()
   4219     # There is no graph to add in eager mode.
   4220     add_to_graph &= not context.executing_eagerly()

~/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py in get_concrete_function(self, *args, **kwargs)
   3149     """
   3150     graph_function = self._get_concrete_function_garbage_collected(
-> 3151         *args, **kwargs)
   3152     graph_function._garbage_collector.release()  # pylint: disable=protected-access
   3153     return graph_function

~/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_garbage_collected(self, *args, **kwargs)
   3114       args, kwargs = None, None
   3115     with self._lock:
-> 3116       graph_function, _ = self._maybe_define_function(args, kwargs)
   3117       seen_names = set()
   3118       captured = object_identity.ObjectIdentitySet(

~/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3461 
   3462           self._function_cache.missed.add(call_context_key)
-> 3463           graph_function = self._create_graph_function(args, kwargs)
   3464           self._function_cache.primary[cache_key] = graph_function
   3465 

~/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3306             arg_names=arg_names,
   3307             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3308             capture_by_value=self._capture_by_value),
   3309         self._function_attributes,
   3310         function_spec=self.function_spec,

~/.local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes, acd_record_initial_resource_uses)
   1005         _, original_func = tf_decorator.unwrap(python_func)
   1006 
-> 1007       func_outputs = python_func(*func_args, **func_kwargs)
   1008 
   1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in wrapped_fn(*args)
   4193           attributes=defun_kwargs)
   4194       def wrapped_fn(*args):  # pylint: disable=missing-docstring
-> 4195         ret = wrapper_helper(*args)
   4196         ret = structure.to_tensor_list(self._output_structure, ret)
   4197         return [ops.convert_to_tensor(t) for t in ret]

~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py in wrapper_helper(*args)
   4123       if not _should_unpack(nested_args):
   4124         nested_args = (nested_args,)
-> 4125       ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
   4126       if _should_pack(ret):
   4127         ret = tuple(ret)

~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    693       except Exception as e:  # pylint:disable=broad-except
    694         if hasattr(e, 'ag_error_metadata'):
--> 695           raise e.ag_error_metadata.to_exception(e)
    696         else:
    697           raise

ValueError: in user code:

    <ipython-input-6-0dff10c5b3ba>:23 make_model_inputs  *
        timestep_grids = occupancy_flow_grids.create_ground_truth_timestep_grids(
    /root/.local/lib/python3.7/site-packages/waymo_open_dataset/utils/occupancy_flow_grids.py:192 create_ground_truth_timestep_grids  *
        current_occupancy = render_func(
    /root/.local/lib/python3.7/site-packages/waymo_open_dataset/utils/occupancy_flow_renderer.py:76 render_occupancy_from_inputs  *
        sampled_points = _sample_and_filter_agent_points(
    /root/.local/lib/python3.7/site-packages/waymo_open_dataset/utils/occupancy_flow_renderer.py:537 _sample_and_filter_agent_points  *
        sampled_points = _sample_agent_points(
    /root/.local/lib/python3.7/site-packages/waymo_open_dataset/utils/occupancy_flow_renderer.py:625 _sample_agent_points  *
        agent_type = _stack_field(inputs, times, 'type')
    /root/.local/lib/python3.7/site-packages/waymo_open_dataset/utils/occupancy_flow_renderer.py:800 _stack_field  *
        fields = tf.broadcast_to(fields[:, :, tf.newaxis, tf.newaxis], x.shape)
    /root/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py:845 broadcast_to  **
        "BroadcastTo", input=input, shape=shape, name=name)
    /root/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:540 _apply_op_helper
        (input_name, err))

    ValueError: Tried to convert 'shape' to a tensor and failed. Error: Cannot convert a partially known TensorShape to a Tensor: (None, 128, 1, 1)

What I understand is that Tensorflow is unable to create a graph that contains the timestep_grids object (and likely the same for true_waypoints and vis_grids). Is there a way to circumvent this issue?

Thank you in advance, Hedi

[rezama]

Can you check the shapes of the input tensors at each stage to see where you start seeing the None shape for the batch dimension?

In case you have already solved it, can you share what fixed it?