def _run_and_report_benchmark(self, **kwargs): start_time_sec = time.time() train_loss, test_loss = distributed_train.main(**kwargs) wall_time_sec = time.time() - start_time_sec extras = {'train_loss': train_loss, 'test_loss': test_loss} self.report_benchmark(wall_time=wall_time_sec, extras=extras)
def _run_and_report_benchmark(self, **kwargs): start_time_sec = time.time() train_loss, test_loss = distributed_train.main(**kwargs) wall_time_sec = time.time() - start_time_sec extras = {'train_loss': train_loss, 'test_loss': test_loss} self.report_benchmark( wall_time=wall_time_sec, extras=extras)
def test_one_epoch_multi_device(self): if tf.test.is_gpu_available(): print('Using 2 virtual GPUs.') device = tf.config.experimental.list_physical_devices('GPU')[0] tf.config.experimental.set_virtual_device_configuration( device, [ tf.config.experimental.VirtualDeviceConfiguration( memory_limit=8192), tf.config.experimental.VirtualDeviceConfiguration( memory_limit=8192) ]) kwargs = utils.get_common_kwargs() kwargs.update({ 'epochs': 1, 'batch_size': 16, 'num_examples': 10, 'embedding_dim': 4, 'enc_units': 4, 'dec_units': 4 }) distributed_train.main(**kwargs)