def testIndexedSlices(self): x = indexed_slices.IndexedSlices( constant_op.constant([1, 2, 3]), constant_op.constant([10, 20, 30])) x_value = indexed_slices.IndexedSlicesValue( np.array([1, 2, 3]), np.array([10, 20, 30]), np.array([100])) self.assertTrue(tensor_util.is_tf_type(x)) self.assertFalse(tensor_util.is_tf_type(x_value))
def testAllReduceSparseVariableLength(self): # One device per process, 2 processes, 2 replicas in total. inputs = [ indexed_slices.IndexedSlicesValue( values=[[1.]], indices=[0], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[2.], [3.], [4.]], indices=[0, 1, 2], dense_shape=[10, 1]), ] expect = indexed_slices.IndexedSlices( values=[[1.], [2.], [3.], [4.]], indices=[0, 0, 1, 2], dense_shape=[10, 1]) self.reduce_and_verify( inputs, expect, self.RunOptions( mode=["func_graph"], # Sparse reduce is not supported in eager. num_processes=2, reduce_op=ReduceOp.SUM))
def testAllReduceSparse(self, num_processes, required_gpus, communication, reduce_op): options = self.RunOptions( mode=["func_graph"], # Sparse reduce is not supported in eager. num_processes=num_processes, gpus_per_process=required_gpus, reduce_op=reduce_op, communication=communication) group_size = options.num_processes * (options.gpus_per_process or 1) inputs_data = [ indexed_slices.IndexedSlicesValue(values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue(values=[[3.], [4.]], indices=[1, 2], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue(values=[[5.], [6.]], indices=[7, 8], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue(values=[[7.], [8.]], indices=[3, 2], dense_shape=[10, 1]), ] inputs = inputs_data[0:group_size] if group_size == 1: expect = indexed_slices.IndexedSlices(values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]) elif group_size == 2: expect = indexed_slices.IndexedSlices(values=[[1.], [2.], [3.], [4.]], indices=[0, 1, 1, 2], dense_shape=[10, 1]) elif group_size == 4: expect = indexed_slices.IndexedSlices( values=[[1.], [2.], [3.], [4.], [5.], [6.], [7.], [8.]], indices=[0, 1, 1, 2, 7, 8, 3, 2], dense_shape=[10, 1]) self.reduce_and_verify(inputs, expect, options)
def _eval_indexed_slices(a): """Converts IndexedSlices to IndexedSlicesValue with numpy indices/values. When eager execution is enabled, converts IndexedSlices to IndexedSlicesValue with numpy indices/values. Args: a: any value. Returns: If a is IndexedSlices and eager execution is enabled, calls numpy() on a's fields. Otherwise returns a unchanged. """ if (isinstance(a, indexed_slices.IndexedSlices) and context.executing_eagerly()): return indexed_slices.IndexedSlicesValue( indices=[x.numpy() for x in a.indices], values=[x.numpy() for x in a.values], dense_shape=a.dense_shape) return a
def testBatchAllReduceSparse(self, num_processes, required_gpus, communication, reduce_op): if required_gpus == 0 and communication == CollectiveCommunication.NCCL: self.skipTest("Skip CPU + NCCL combination") if num_processes == 2 and communication == CollectiveCommunication.NCCL: self.skipTest("Skip NCCL + 2 processes combination. NCCL requires " "physical GPUs for every process.") options = self.RunOptions( mode=["func_graph"], # Sparse reduce is not supported in eager. num_processes=num_processes, gpus_per_process=required_gpus, reduce_op=reduce_op, communication=communication) group_size = options.num_processes * (options.gpus_per_process or 1) inputs_data = ([ indexed_slices.IndexedSlicesValue( values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[3.], [4.]], indices=[1, 2], dense_shape=[5, 1]) ], [ indexed_slices.IndexedSlicesValue( values=[[5.], [6.]], indices=[1, 2], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[7.], [8.]], indices=[0, 1], dense_shape=[5, 1]) ], [ indexed_slices.IndexedSlicesValue( values=[[9.], [10.]], indices=[3, 4], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[11.], [12.]], indices=[3, 4], dense_shape=[5, 1]) ], [ indexed_slices.IndexedSlicesValue( values=[[13.], [14.]], indices=[8, 9], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[15.], [16.]], indices=[3, 4], dense_shape=[5, 1]) ]) inputs = inputs_data[0:group_size] if group_size == 1: expect = [ indexed_slices.IndexedSlices( values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]), indexed_slices.IndexedSlicesValue( values=[[3.], [4.]], indices=[1, 2], dense_shape=[5, 1]) ] if group_size == 2: expect = [ indexed_slices.IndexedSlices( values=[[1.], [2.], [5.], [6.]], indices=[0, 1, 1, 2], dense_shape=[10, 1]), indexed_slices.IndexedSlices( values=[[3.], [4.], [7.], [8.]], indices=[1, 2, 3, 4], dense_shape=[5, 1]) ] elif group_size == 4: expect = [ indexed_slices.IndexedSlices( values=[[1.], [2.], [5.], [6.], [9.], [10.], [13.], [14.]], indices=[0, 1, 1, 2, 3, 4, 8, 9], dense_shape=[10, 1]), indexed_slices.IndexedSlices( values=[[3.], [4.], [7.], [8.], [11.], [12.], [15.], [16.]], indices=[1, 2, 0, 1, 3, 4, 3, 4], dense_shape=[5, 2]) ] self.batch_reduce_and_verify(inputs, expect, options)