def pack_small_tensors(tower_grads, max_bytes=0):
    """Concatenate gradients together more intelligently.

  Does binpacking
  Args:
    tower_grads: List of lists of (gradient, variable) tuples.
    max_bytes: Int giving max number of bytes in a tensor that
      may be considered small.
  """
    assert max_bytes >= 0
    orig_grads = [g for g, _ in tower_grads[0]]
    # Check to make sure sizes are accurate; not entirely important
    assert all(g.dtype == tf.float32 for g in orig_grads)
    sizes = [4 * g.shape.num_elements() for g in orig_grads]
    print_stats(sizes)
    small_ranges = []
    large_indices = []
    new_sizes = []

    def end_interval(indices, small_ranges, large_indices):
        if len(indices) > 1:
            small_ranges.insert(0, [indices[0], indices[-1]])
        else:
            large_indices.insert(0, indices[0])

    cur_range = []
    cur_size = 0
    for i, s in reversed(list(enumerate(sizes))):
        if cur_size > max_bytes:
            end_interval(cur_range, small_ranges, large_indices)
            new_sizes.insert(0, cur_size)
            cur_range = []
            cur_size = 0
        cur_range.insert(0, i)
        cur_size += s
    end_interval(cur_range, small_ranges, large_indices)
    new_sizes.insert(0, cur_size)

    print_stats(new_sizes)
    num_gv = len(orig_grads)
    packing = {}
    if len(small_ranges):
        new_tower_grads = []
        for dev_idx, gv_list in enumerate(tower_grads):
            assert len(gv_list) == num_gv, (
                "Possible cause: "
                "Networks constructed on different workers "
                "don't have the same number of variables. "
                "If you use tf.GraphKeys or tf.global_variables() "
                "with multiple graphs per worker during network "
                "construction, you need to use "
                "appropriate scopes, see "
                "https://github.com/ray-project/ray/issues/3136")
            new_gv_list = []
            for r in small_ranges:
                key = "%d:%d" % (dev_idx, len(new_gv_list))
                new_gv_list.append((pack_range(key, packing, gv_list, r),
                                    "packing_var_placeholder"))
            for i in large_indices:
                new_gv_list.append(gv_list[i])
            new_tower_grads.append(new_gv_list)
        return new_tower_grads, packing
    else:
        return tower_grads, None
def pack_small_tensors ( tower_grads , max_bytes = 0 ) : assert max_bytes >= 0 orig_grads = [ g for g , _ in tower_grads [ 0 ] ] assert all ( g . dtype == tf . float32 for g in orig_grads ) sizes = [ 4 * g . shape . num_elements ( ) for g in orig_grads ] print_stats ( sizes ) small_ranges = [ ] large_indices = [ ] new_sizes = [ ] def end_interval ( indices , small_ranges , large_indices ) : if len ( indices ) > 1 : small_ranges . insert ( 0 , [ indices [ 0 ] , indices [ - 1 ] ] ) else : large_indices . insert ( 0 , indices [ 0 ] ) cur_range = [ ] cur_size = 0 for i , s in reversed ( list ( enumerate ( sizes ) ) ) : if cur_size > max_bytes : end_interval ( cur_range , small_ranges , large_indices ) new_sizes . insert ( 0 , cur_size ) cur_range = [ ] cur_size = 0 cur_range . insert ( 0 , i ) cur_size += s end_interval ( cur_range , small_ranges , large_indices ) new_sizes . insert ( 0 , cur_size ) print_stats ( new_sizes ) num_gv = len ( orig_grads ) packing = { } if len ( small_ranges ) : new_tower_grads = [ ] for dev_idx , gv_list in enumerate ( tower_grads ) : assert len ( gv_list ) == num_gv , ( "Possible cause: " "Networks constructed on different workers " "don't have the same number of variables. " "If you use tf.GraphKeys or tf.global_variables() " "with multiple graphs per worker during network " "construction, you need to use " "appropriate scopes, see " "https://github.com/ray-project/ray/issues/3136" ) new_gv_list = [ ] for r in small_ranges : key = "%d:%d" % ( dev_idx , len ( new_gv_list ) ) new_gv_list . append ( ( pack_range ( key , packing , gv_list , r ) , "packing_var_placeholder" ) ) for i in large_indices : new_gv_list . append ( gv_list [ i ] ) new_tower_grads . append ( new_gv_list ) return new_tower_grads , packing else : return tower_grads , None
Concatenate gradients together more intelligently.

  Does binpacking
  Args:
    tower_grads: List of lists of (gradient, variable) tuples.
    max_bytes: Int giving max number of bytes in a tensor that
      may be considered small.