Python create_device_ids_array Exemples

Exemple #1

Fichier : optimizers_test.py Projet : ttigong/keras

 def setUp(self):
     super(OptimizersTest, self).setUp()
     global_ids = test_util.create_device_ids_array((2, 2))
     local_device_ids = np.ravel(global_ids).tolist()
     mesh_dict = {
         'CPU':
         dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
                      test_util.create_device_list((2, 2), 'CPU'))
     }
     self.mesh = self.configTestMesh(mesh_dict)

Exemple #2

 def setUp(self):
     super(UtilsTest, self).setUp()
     global_ids = test_util.create_device_ids_array((2, 2))
     local_device_ids = np.ravel(global_ids).tolist()
     mesh_dict = {
         'CPU':
         dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
                      test_util.create_device_list((2, 2), 'CPU'))
     }
     self.mesh = self.configTestMesh(mesh_dict)
     self.layout = dtensor.Layout.replicated(self.mesh, rank=1)

Exemple #3

Fichier : layers_test.py Projet : vishalbelsare/keras

 def setUp(self):
     super(LayersTest, self).setUp()
     backend.enable_tf_random_generator()
     tf_utils.set_random_seed(1337)
     global_ids = test_util.create_device_ids_array((2, 2))
     local_device_ids = np.ravel(global_ids).tolist()
     mesh_dict = {
         'CPU':
         dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
                      test_util.create_device_list((2, 2), 'CPU'))
     }
     self.mesh = self.configTestMesh(mesh_dict)

Exemple #4

 def setUp(self):
   super(LayersTest, self).setUp()
   tf_utils.set_random_seed(1337)
   global_ids = test_util.create_device_ids_array((2, 2))
   local_device_ids = np.ravel(global_ids).tolist()
   mesh_dict = {
       'CPU':
           dtensor.Mesh(['X', 'Y'], global_ids,
                        local_device_ids,
                        test_util.create_device_list((2, 2), 'CPU'))
   }
   self.mesh = self.configTestMesh(mesh_dict)
   self.layout_4d = dtensor.Layout.replicated(self.mesh, rank=4)
   self.layout_3d = dtensor.Layout.replicated(self.mesh, rank=3)
   self.layout_2d = dtensor.Layout.replicated(self.mesh, rank=2)
   self.layout_1d = dtensor.Layout.replicated(self.mesh, rank=1)

Exemple #5

Fichier : layout_map_test.py Projet : Alan-love/keras

    def setUp(self):
        super(LayoutMapTest, self).setUp()
        backend.enable_tf_random_generator()
        tf_utils.set_random_seed(1337)
        global_ids = test_util.create_device_ids_array((2, 2))
        local_device_ids = np.ravel(global_ids).tolist()
        mesh_dict = {
            'CPU':
            dtensor.Mesh(['X', 'Y'], global_ids, local_device_ids,
                         test_util.create_device_list((2, 2), 'CPU'))
        }
        self.mesh = self.configTestMesh(mesh_dict)
        self.layout_2d = dtensor.Layout.replicated(self.mesh, rank=2)
        self.layout_1d = dtensor.Layout.replicated(self.mesh, rank=1)

        self.sharded_2d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=2)
        self.sharded_1d = dtensor.Layout.batch_sharded(self.mesh, 'X', rank=1)