def get_final_graph_state(G):
'''
Build a graph that only uses the final time step of the node vectors and the edges
'''
all_nodes = G['nodes']
last_nodes = {}
for nk, nodes in all_nodes.items():
node_attrs = {}
for attr, vals in nodes.items():
node_attrs[attr] = vals[:,-1:]
num_times = vals.shape.as_list()[1]
last_nodes[nk] = propdict(node_attrs)
all_edges = G['edges']
final_edges = {}
for ek, edges in all_edges.items():
edge_attrs = {}
last_time_inds = tf.cast(tf.where(tf.equal(edges['idx'][:,1], tf.ones_like(edges['idx'][:,1]) * (num_times-1))), tf.int32) # [num_last_edges,1]
for attr, vals in edges.items():
if attr == 'layer':
edge_attrs[attr] = vals
elif attr == 'idx':
last_edges = tf.gather_nd(vals, last_time_inds) # [num_last_edges, 4]
last_edges = tf.concat([last_edges[:,0:1], tf.zeros_like(last_edges[:,0:1]), last_edges[:,2:]], axis=-1) # all time inds are 0
edge_attrs[attr] = last_edges
else:
edge_attrs[attr] = tf.gather_nd(vals, last_time_inds)
final_edges[ek] = propdict(edge_attrs)
return Graph(nodes=last_nodes, edges=final_edges)
def vectors_to_attrs(self, node_vectors, attribute_dims, attribute_preprocs={}):
'''
Converts nodes in vector format to attribute format (i.e. a dict of attributes)
node_vectors: [B,T,N,D] <tf.float32>
attribute_dims: <dict> of (attr, [d0,d1]) pairs. By convention, d1 == 0 sets d1 <- D
attribute_preprocs: <dict> of (attr, func) pairs to apply to each attr.
returns:
attrs: <dict> of (attr, [B,T,N, attribute_dims[attr][1] - attribute_dims[attr][0]] <tf.float32>) processed attrs
'''
B,T,N,D = node_vectors.shape.as_list()
attr_outputs = {}
attr_dims = {}
for attr, dims in attribute_dims.items():
d0, d1 = [(D+d if d<0 else d) for d in dims]
d1 = D if d1==0 else d1
func = attribute_preprocs.get(attr, tf.identity) or tf.identity
attr_func = lambda tensor: tf.identity(func(tensor), name=(attr+'_preproc'))
attr_out = attr_func(node_vectors[...,d0:d1])
if attr in self.stop_gradient_on_attributes:
print("stopping gradient on %s" % attr)
attr_out = tf.stop_gradient(attr_out)
attr_outputs[attr] = attr_out
attr_dims[attr] = [d0, d1]
return propdict(attr_outputs), attr_dims
def all_to_all_edges_from_nodes(graph, edge_key, node_key=None, self_edges=False):
'''
Adds 'all_to_all' edges to a graph at a level specified by the edge or node key
'''
if node_key is None:
assert 'all_to_all' in edge_key
level = edge_key.split('_')[-1] # str(int) of level
node_key = 'nodes_level_%s' % level
assert node_key in graph.nodes.keys()
valid_nodes = graph.nodes[node_key]['valid'] # [B,T,N,1]
B,T,N,_ = valid_nodes.shape.as_list()
valid_edges = (valid_nodes * tf.transpose(valid_nodes, [0,1,3,2])) > 0.5 # bool
if not self_edges:
valid_edges = tf.logical_and(valid_edges, tf.logical_not(tf.eye(N, batch_shape=[B,T], dtype=tf.bool)))
valid_edges = tf.cast(tf.where(valid_edges), tf.int32) # [?,4]
level = node_key.split('_')[-1]
edge_key = 'all_to_all_edges_%s_to_%s' % (level, level)
graph.edges[edge_key] = propdict({
'idx': valid_edges,
'layer': np.array([node_key, node_key]).reshape([1,2])})
return graph
def copy_graph(graph):
nodes = graph.nodes
nodes_copy = {
nk: propdict({
ak: tf.identity(nodes[nk][ak], name=('%s_copy'%ak)) for ak in nodes[nk].keys()
}) for nk in nodes.keys()
}
edges = graph.edges
edges_copy = {
ek: propdict({
ak: (tf.identity(edges[ek][ak], name=('%s_copy'%ak)) if ak != 'layer' else np.copy(edges[ek][ak]))
for ak in edges[ek].keys()
}) for ek in edges.keys()
}
graph_copy = Graph(nodes=nodes_copy, edges=edges_copy)
return graph_copy
def across_parent_edges_from_nodes(graph, edge_key):
assert 'across' in edge_key
lev1, lev2 = edge_key.split('_')[-3::2]
assert lev1 == lev2, "edge key must be connecting same level but levels are %s, %s" % (lev1, lev2)
level = lev1
node_key = 'nodes_level_%s' % level
assert node_key in graph.nodes.keys()
c2p_key = 'child_to_parent_edges_%s_to_%s' % (level, str(int(level)+1))
assert c2p_key in graph.edges.keys(), "c2p_key: %s" % (c2p_key)
c2p_edges = graph.edges[c2p_key]['idx'] # [?,4] inds
valid_nodes = tf.cast(graph.nodes[node_key]['valid'], tf.bool) # [B,T,N,1]
parents = tf.scatter_nd(indices=c2p_edges[:,0:3], updates=c2p_edges[:,3:4], shape=valid_nodes.shape) # [B,T,N,1]
diff_parents = tf.logical_not(tf.equal(parents, tf.transpose(parents, [0,1,3,2]))) # [B,T,N,N]
diff_parents = tf.logical_and(diff_parents, tf.logical_and(valid_nodes, tf.transpose(valid_nodes, [0,1,3,2]))) # [B,T,N,N]
graph.edges[edge_key] = propdict({
'idx': tf.cast(tf.where(diff_parents), tf.int32),
'layer': np.array([node_key, node_key]).reshape([1,2])})
return graph
def build_psg(self, base_tensor, scope='GraphBuild', **kwargs):
### TODO ###
with tf.compat.v1.variable_scope(scope):
nodes_level_0 = tf.reshape(
base_tensor, [-1, self.Tb, self.Hb * self.Wb, self.Cb])
self.psg = Graph(edges={},
nodes={
'level0':
propdict(
self.baseDims.get_tensor_from_attrs(
nodes_level_0,
self.baseDims.keys(),
concat=False))
})
self.psg.dims = {'level0': self.baseDims.copy(suffix='')}
self.psg.spatial = {
'level0': utils.inds_image(self.B, self.Tb, [self.Hb, self.Wb])
}
return self.psg
def within_parent_edges_from_nodes(graph, edge_key, self_edges=False):
assert 'within_parent' in edge_key
lev1, lev2 = edge_key.split('_')[-3::2]
assert lev1 == lev2
level = lev1
c2p_key = 'child_to_parent_edges_%s_to_%s' % (level, str(int(level)+1))
assert c2p_key in graph.edges.keys()
c2p_edges = graph.edges[c2p_key]['idx'] # [?,4]
node_key = graph.edges[c2p_key]['layer'].squeeze()[0]
valid_nodes = tf.cast(graph.nodes[node_key]['valid'], tf.bool) # [B,T,N,1]
B,T,N,_ = valid_nodes.shape.as_list()
parents = tf.scatter_nd(indices=c2p_edges[:,0:3], updates=c2p_edges[:,3:4], shape=valid_nodes.shape) # [B,T,N,1]
same_parents = tf.equal(parents, tf.transpose(parents, [0,1,3,2]))
if not self_edges:
same_parents = tf.logical_and(same_parents, tf.logical_not(tf.eye(N, batch_shape=[B,T], dtype=tf.bool)))
same_parents = tf.logical_and(same_parents, tf.logical_and(valid_nodes, tf.transpose(valid_nodes, [0,1,3,2]))) # [B,T,N,N]
graph.edges[edge_key] = propdict({
'idx': tf.cast(tf.where(same_parents), tf.int32),
'layer': np.array([node_key, node_key]).reshape([1,2])
})
return graph
def construct_graph(self,
nodes_list,
nodes_to_parents_list=None,
object_level_kwargs={},
**kwargs
):
'''
Builds the graph in the format used for dynamics models
nodes_list: len(self.num_levels) <list> of tf.Tensors [B,T,N_level,D_level] <tf.float32>
nodes_to_parents_list: len(self.num_levels - 1) <list> of tf.Tensors [B,T,N_level] <tf.int32> indices in range [0, N_(level+1) )
'''
self.num_levels = len(nodes_list) + int(self.build_object_nodes)
if isinstance(nodes_list, OrderedDict):
nodes_list = nodes_list.values()
if isinstance(nodes_to_parents_list, OrderedDict):
nodes_to_parents_list = nodes_to_parents_list.values()
# compute edges up and down the hierarchy (effectively pointers to nodes)
if nodes_to_parents_list is None:
# TODO: find nearest neighbors
raise NotImplementedError("This should just be a nearest neighbors function")
# build object nodes
if self.build_object_nodes:
object_nodes, n2obj_inds = self.add_object_level(nodes_list[-1], **object_level_kwargs)
nodes_list.append(object_nodes)
nodes_to_parents_list.append(n2obj_inds)
# convert the nodes to attrs
all_nodes = {}
attribute_dims_dict = {}
attribute_preprocs_dict = {}
for level, nodes in enumerate(nodes_list):
node_key = 'nodes_level_' + str(level)
if self.Dims is not None:
Dims = self.Dims[level]
node_attrs = Dims.get_tensor_from_attrs(nodes, Dims.keys(), postproc=True, concat=False,
stop_gradient={a:True for a in self.stop_gradient_on_attributes})
node_attrs = propdict(node_attrs)
new_attr_dims = copy.deepcopy({attr:Dims[attr][:2] for attr in Dims.keys()})
else:
node_attrs, new_attr_dims = self.vectors_to_attrs(nodes, self.attribute_dims[level], self.attribute_preprocs[level])
all_nodes[node_key] = node_attrs
# alter the attribute dims and preprocs
attribute_dims_dict[node_key] = new_attr_dims
attribute_preprocs_dict[node_key] = self.attribute_preprocs[level]
self.attribute_dims = attribute_dims_dict
self.attribute_preprocs = attribute_preprocs_dict
all_edges = {}
for level, n2p_inds in enumerate(nodes_to_parents_list):
assert n2p_inds.shape.as_list() == nodes_list[level].shape.as_list()[:-1], (n2p_inds.shape.as_list(), nodes_list[level].shape.as_list())
assert n2p_inds.dtype == tf.int32
c2p_key = 'child_to_parent_edges_%d_to_%d' % (level, level+1)
p2c_key = 'parent_to_child_edges_%d_to_%d' % (level+1, level)
valid_nodes = tf.cast(all_nodes['nodes_level_'+str(level)]['valid'][...,0], tf.bool) # [B,T,N] in {False, True}
# add edges to dict
c2p_rect = graphical.add_batch_time_node_index(n2p_inds) # [B,T,N,4] tensor of (b_ind, t_ind, child_ind, parent_ind)
c2p_list = tf.gather_nd(c2p_rect, tf.cast(tf.where(valid_nodes), tf.int32)) # [?, 4]
all_edges[c2p_key] = propdict({
'layer': np.array(['nodes_level_'+str(lev) for lev in [level, level+1]]).reshape([1,2]),
'idx': c2p_list})
p2c_list = tf.concat([c2p_list[:,0:2], c2p_list[:,3:4], c2p_list[:,2:3]], axis=-1)
all_edges[p2c_key] = propdict({
'layer': np.array(['nodes_level_'+str(lev) for lev in [level+1, level]]).reshape([1,2]),
'idx': p2c_list})
# now build graph
self.graph = Graph(nodes=all_nodes, edges=all_edges)
