def testUnpackSmallTensors(self): packing = {'0:0': allreduce.GradPackTuple(indices=range(2), vars=['v_0_0', 'v_0_1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]), '0:1': allreduce.GradPackTuple(indices=range(3, 5), vars=['v_0_3', 'v_0_4'], shapes=[tf.TensorShape([3, 3,]), tf.TensorShape([3, 3,])]), '1:0': allreduce.GradPackTuple(indices=range(2), vars=['v_1_0', 'v_1_1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]), '1:1': allreduce.GradPackTuple(indices=range(3, 5), vars=['v_1_3', 'v_1_4'], shapes=[tf.TensorShape([3, 3,]), tf.TensorShape([3, 3,])])} t0 = tf.constant([0, 1, 2, 3, 4, 5, 6, 7], dtype=tf.float32) t1 = tf.constant([17, 17], dtype=tf.float32) t2 = tf.constant([0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=tf.float32) t3 = tf.constant([0], dtype=tf.float32) tower_grads = [] for d in range(0, 2): one_tower = [(t0, 'packing_var_placeholder'), (t2, 'packing_var_placeholder'), (t1, 'v_%d_2' % d), (t3, 'v_%d_5' %d)] tower_grads.append(one_tower) new_tower_grads = allreduce.unpack_small_tensors(tower_grads, packing) self.assertEqual(2, len(new_tower_grads)) for d, tg in enumerate(new_tower_grads): self.assertEqual(6, len(tg)) self.assertEqual('v_%d_0' % d, tg[0][1]) self.assertEqual('v_%d_1' % d, tg[1][1]) self.assertEqual('v_%d_2' % d, tg[2][1]) self.assertEqual('v_%d_3' % d, tg[3][1]) self.assertEqual('v_%d_4' % d, tg[4][1]) self.assertEqual('v_%d_5' % d, tg[5][1]) self.assertEqual(1, tg[0][0].shape.ndims) self.assertEqual(4, tg[0][0].shape.dims[0]) self.assertEqual(1, tg[1][0].shape.ndims) self.assertEqual(4, tg[1][0].shape.dims[0]) self.assertEqual(1, tg[2][0].shape.ndims) self.assertEqual(2, tg[2][0].shape.dims[0]) self.assertEqual(2, tg[3][0].shape.ndims) self.assertEqual(3, tg[3][0].shape.dims[0]) self.assertEqual(3, tg[3][0].shape.dims[1]) self.assertEqual(2, tg[4][0].shape.ndims) self.assertEqual(3, tg[4][0].shape.dims[0]) self.assertEqual(3, tg[4][0].shape.dims[1]) self.assertEqual(1, tg[5][0].shape.ndims) self.assertEqual(1, tg[5][0].shape.dims[0])
def testUnpackGradTuple(self): packing = { '0:0': allreduce.GradPackTuple(indices=range(4), vars=['v0', 'v1', 'v2', 'v3'], shapes=[ tf.TensorShape([4]), tf.TensorShape([4]), tf.TensorShape([3, 3]), tf.TensorShape([3, 3]) ]) } tc = tf.constant([ 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8 ], dtype=tf.float32) packed_gv = [tc, 'packing_var_placeholder'] gv = allreduce.unpack_grad_tuple(packed_gv, packing['0:0']) self.assertEqual(4, len(gv)) self.assertEqual('v0', gv[0][1]) self.assertEqual('v1', gv[1][1]) self.assertEqual('v2', gv[2][1]) self.assertEqual('v3', gv[3][1]) self.assertEqual(1, gv[0][0].shape.ndims) self.assertEqual(4, gv[0][0].shape.dims[0]) self.assertEqual(1, gv[1][0].shape.ndims) self.assertEqual(4, gv[1][0].shape.dims[0]) self.assertEqual(2, gv[2][0].shape.ndims) self.assertEqual(3, gv[2][0].shape.dims[0]) self.assertEqual(3, gv[2][0].shape.dims[1])
def testPackRange(self): packing = {} t0 = tf.constant([0, 1, 2, 3], dtype=tf.float32) t1 = tf.constant([4, 5, 6, 7], dtype=tf.float32) gv = [(t0, 'v0'), (t1, 'v1')] new_t = allreduce.pack_range('0:0', packing, gv, [0, 1]) self.assertEqual(1, new_t.shape.ndims) self.assertEqual(8, new_t.shape.dims[0]) self.assertEqual( packing, { '0:0': allreduce.GradPackTuple( indices=range(2), vars=['v0', 'v1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]) }) t2 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=tf.float32) t3 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=tf.float32) gv = [(t0, 'v0'), (t1, 'v1'), (t2, 'v2'), (t3, 'v3')] packing = {} new_t = allreduce.pack_range('1:0', packing, gv, [0, 3]) self.assertEqual(1, new_t.shape.ndims) self.assertEqual(26, new_t.shape.dims[0]) self.assertEqual( packing, { '1:0': allreduce.GradPackTuple( indices=range(4), vars=['v0', 'v1', 'v2', 'v3'], shapes=[ tf.TensorShape([4]), tf.TensorShape([4]), tf.TensorShape([3, 3]), tf.TensorShape([3, 3]) ]) })
def testPackSmallTensors(self): t0 = tf.constant([0, 1, 2, 3], dtype=tf.float32) t1 = tf.constant([4, 5, 6, 7], dtype=tf.float32) t2 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=tf.float32) t3 = tf.constant([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=tf.float32) tower_grads = [] for d in range(0, 3): gv = [(t0, 'v_%d_0' % d), (t1, 'v_%d_1' % d), (t2, 'v_%d_2' % d), (t3, 'v_%d_3' % d)] tower_grads.append(gv) # 1) Set the size limit so small that nothing gets concatenated. new_tower_grads, packing = allreduce.pack_small_tensors(tower_grads, max_bytes=12, max_group=10) self.assertEqual(tower_grads, new_tower_grads) self.assertTrue(packing is None) # 2) Set the size limit so only the first two tensors get concatenated new_tower_grads, packing = allreduce.pack_small_tensors( tower_grads, max_bytes=16, # 16 bytes == 4 elements max_group=10) self.assertEqual(3, len(new_tower_grads)) self.assertEqual(4, len(tower_grads[0])) first_tower = new_tower_grads[0] self.assertEqual(3, len(first_tower)) self.assertEqual(1, first_tower[0][0].shape.ndims) self.assertEqual(8, first_tower[0][0].shape.dims[0]) self.assertEqual( packing, { '0:0': allreduce.GradPackTuple( indices=range(2), vars=['v_0_0', 'v_0_1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]), '1:0': allreduce.GradPackTuple( indices=range(2), vars=['v_1_0', 'v_1_1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]), '2:0': allreduce.GradPackTuple( indices=range(2), vars=['v_2_0', 'v_2_1'], shapes=[tf.TensorShape([4]), tf.TensorShape([4])]) }) # 3) Set the size limit so all tensors get concatenated new_tower_grads, packing = allreduce.pack_small_tensors( tower_grads, max_bytes=256, # bytes = 64 elements max_group=10) self.assertEqual(3, len(new_tower_grads)) self.assertEqual(4, len(tower_grads[0])) self.assertEqual(1, len(new_tower_grads[0])) first_tower = new_tower_grads[0] self.assertEqual(1, first_tower[0][0].shape.ndims) self.assertEqual(26, first_tower[0][0].shape.dims[0]) self.assertEqual( packing, { '0:0': allreduce.GradPackTuple( indices=range(4), vars=['v_0_0', 'v_0_1', 'v_0_2', 'v_0_3'], shapes=[ tf.TensorShape([4]), tf.TensorShape([4]), tf.TensorShape([ 3, 3, ]), tf.TensorShape([ 3, 3, ]) ]), '1:0': allreduce.GradPackTuple( indices=range(4), vars=['v_1_0', 'v_1_1', 'v_1_2', 'v_1_3'], shapes=[ tf.TensorShape([4]), tf.TensorShape([4]), tf.TensorShape([ 3, 3, ]), tf.TensorShape([ 3, 3, ]) ]), '2:0': allreduce.GradPackTuple( indices=range(4), vars=['v_2_0', 'v_2_1', 'v_2_2', 'v_2_3'], shapes=[ tf.TensorShape([4]), tf.TensorShape([4]), tf.TensorShape([ 3, 3, ]), tf.TensorShape([ 3, 3, ]) ]) })