Reputation: 178
While loading a pretrained model with pytorch via
model.load_state_dict(torch.load(MODEL_PATH))
the console is flooded with output containing information about the model (which is annoying). Afaik there is no verbosity option, neither in model.load_state_dict
nor in torch.load
. Please let me know if I overlooked this parameter.
However, this led me to the question if there is any general way to force a function to be non verbose. Maybe something like:
with os.nonverbose():
model.load_state_dict(torch.load(MODEL_PATH))
Any ideas?
Upvotes: 0
Views: 2223
Reputation: 3301
As @triplee commented, most libraries would use Python logging, which can be modified extensively. I haven't worked with PyTorch before but from https://github.com/pytorch/vision/issues/330 it looks like it's actaully using print
(which is horrible if it is but okay).
In general, you can, however, suppress the stdout output of anything by redirecting stdout. Here's a good link https://wrongsideofmemphis.com/2010/03/01/store-standard-output-on-a-variable-in-python/
In your question you ask if this can be done with a context manager. I don't see why not and it seems appropriate given that you want to reset stdout after the function call. Something like this:
from io import StringIO # Python3
import sys
class SilencedStdOut:
def __enter__(self):
self.old_stdout = sys.stdout
self.result = StringIO()
sys.stdout = self.result
def __exit__(self, *args, **kwargs):
sys.stdout = self.old_stdout
result_string = self.result.getvalue() # use if you want or discard.
If you only ever want to supress a single function and never a block of code, however, then a decorator should work fine too.
Upvotes: 1