Переглянути джерело

block printing while importing paddleclas

will-jl944 4 роки тому
батько
коміт
196e2a5cb0

+ 5 - 3
dygraph/paddlex/cv/models/classifier.py

@@ -21,11 +21,13 @@ import paddle
 from paddle import to_tensor
 import paddle.nn.functional as F
 from paddle.static import InputSpec
-from paddlex.utils import logging, TrainingStats
+from paddlex.utils import logging, TrainingStats, DisablePrint
 from paddlex.cv.models.base import BaseModel
 from paddlex.cv.transforms import arrange_transforms
-from PaddleClas.ppcls.modeling import architectures
-from PaddleClas.ppcls.modeling.loss import CELoss
+
+with DisablePrint():
+    from PaddleClas.ppcls.modeling import architectures
+    from PaddleClas.ppcls.modeling.loss import CELoss
 
 __all__ = [
     "ResNet18", "ResNet34", "ResNet50", "ResNet101", "ResNet152",

+ 2 - 1
dygraph/paddlex/utils/__init__.py

@@ -15,7 +15,8 @@
 from . import logging
 from . import utils
 from .utils import (seconds_to_hms, get_encoding, get_single_card_bs, dict2str,
-                    EarlyStop, path_normalization, is_pic, MyEncoder)
+                    EarlyStop, path_normalization, is_pic, MyEncoder,
+                    DisablePrint)
 from .checkpoint import get_pretrain_weights, load_pretrain_weights
 from .env import get_environ_info, get_num_workers, init_parallel_env
 from .download import download_and_decompress

+ 12 - 0
dygraph/paddlex/utils/utils.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+import os
 import math
 import chardet
 import json
@@ -126,3 +128,13 @@ class EarlyStop:
                 self.counter = 0
                 self.score = current_score
                 return False
+
+
+class DisablePrint(object):
+    def __enter__(self):
+        self._original_stdout = sys.stdout
+        sys.stdout = open(os.devnull, 'w')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout.close()
+        sys.stdout = self._original_stdout