Python configure_logical_cpus Examples

Example #1

File: learner_test.py Project: SamuelMarks/agents

    def setUp(self):
        super(LearnerTest, self).setUp()

        # Define 4 logical CPUs on the first physical one.
        dist_test_utils.configure_logical_cpus()

        devices_gpu = tf.config.list_physical_devices('GPU')
        # If there are GPU devices:
        if devices_gpu:
            tf.config.experimental.set_virtual_device_configuration(
                devices_gpu[0], [
                    tf.config.LogicalDeviceConfiguration(memory_limit=10),
                    tf.config.LogicalDeviceConfiguration(memory_limit=10),
                    tf.config.LogicalDeviceConfiguration(memory_limit=10),
                    tf.config.LogicalDeviceConfiguration(memory_limit=10)
                ])

Example #2

 def setUp(self):
     super(TrainUtilsTest, self).setUp()
     dist_test_utils.configure_logical_cpus()