diff --git a/official/utils/logs/logger.py b/official/utils/logs/logger.py index fb4b8870f23..041111ec17a 100644 --- a/official/utils/logs/logger.py +++ b/official/utils/logs/logger.py @@ -213,15 +213,18 @@ def _collect_cpu_info(run_info): cpu_info["num_cores"] = multiprocessing.cpu_count() - # Note: cpuinfo is not installed in the TensorFlow OSS tree. - # It is installable via pip. - import cpuinfo # pylint: disable=g-import-not-at-top + try: + # Note: cpuinfo is not installed in the TensorFlow OSS tree. + # It is installable via pip. + import cpuinfo # pylint: disable=g-import-not-at-top - info = cpuinfo.get_cpu_info() - cpu_info["cpu_info"] = info["brand"] - cpu_info["mhz_per_cpu"] = info["hz_advertised_raw"][0] / 1.0e6 + info = cpuinfo.get_cpu_info() + cpu_info["cpu_info"] = info["brand"] + cpu_info["mhz_per_cpu"] = info["hz_advertised_raw"][0] / 1.0e6 - run_info["machine_config"]["cpu_info"] = cpu_info + run_info["machine_config"]["cpu_info"] = cpu_info + except ImportError: + tf.logging.warn("'cpuinfo' not imported. CPU info will not be logged.") def _collect_gpu_info(run_info): @@ -243,12 +246,15 @@ def _collect_gpu_info(run_info): def _collect_memory_info(run_info): - # Note: psutil is not installed in the TensorFlow OSS tree. - # It is installable via pip. - import psutil # pylint: disable=g-import-not-at-top - vmem = psutil.virtual_memory() - run_info["machine_config"]["memory_total"] = vmem.total - run_info["machine_config"]["memory_available"] = vmem.available + try: + # Note: psutil is not installed in the TensorFlow OSS tree. + # It is installable via pip. + import psutil # pylint: disable=g-import-not-at-top + vmem = psutil.virtual_memory() + run_info["machine_config"]["memory_total"] = vmem.total + run_info["machine_config"]["memory_available"] = vmem.available + except ImportError: + tf.logging.warn("'psutil' not imported. Memory info will not be logged.") def _parse_gpu_model(physical_device_desc):