Skip to content

Commit 03781c7

Browse files
authored
Update the importing logic for cpuinfo and psutil. (#3781)
* Update the importing logic for cpuinfo and psutil. Those two libs are usually not installed by default, and we should not force people to install them if they just want to run resnet. * Add pylint warning suppression.
1 parent 092c966 commit 03781c7

File tree

1 file changed

+7
-7
lines changed

1 file changed

+7
-7
lines changed

official/utils/logging/logger.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,6 @@
2828
import numbers
2929
import os
3030

31-
# pylint: disable=g-bad-import-order
32-
# Note: cpuinfo and psutil are not installed in the TensorFlow OSS tree.
33-
# They are installable via pip.
34-
import cpuinfo
35-
import psutil
36-
# pylint: enable=g-bad-import-order
37-
3831
import tensorflow as tf
3932
from tensorflow.python.client import device_lib
4033

@@ -149,6 +142,10 @@ def _collect_cpu_info(run_info):
149142

150143
cpu_info["num_cores"] = multiprocessing.cpu_count()
151144

145+
# Note: cpuinfo is not installed in the TensorFlow OSS tree.
146+
# It is installable via pip.
147+
import cpuinfo # pylint: disable=g-import-not-at-top
148+
152149
info = cpuinfo.get_cpu_info()
153150
cpu_info["cpu_info"] = info["brand"]
154151
cpu_info["mhz_per_cpu"] = info["hz_advertised_raw"][0] / 1.0e6
@@ -175,6 +172,9 @@ def _collect_gpu_info(run_info):
175172

176173

177174
def _collect_memory_info(run_info):
175+
# Note: psutil is not installed in the TensorFlow OSS tree.
176+
# It is installable via pip.
177+
import psutil # pylint: disable=g-import-not-at-top
178178
vmem = psutil.virtual_memory()
179179
run_info["memory_total"] = vmem.total
180180
run_info["memory_available"] = vmem.available

0 commit comments

Comments
 (0)