0

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?

justanyphil
  • 138
  • 8

1 Answers1

1

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.

Neil
  • 3,020
  • 4
  • 25
  • 48
  • 1
    This is a brilliant solution, I'd never thought of overwriting stdouts in python. – smassey Jan 07 '21 at 08:31
  • 3
    From the standard lib, there's also [contextlib.redirect_stdout](https://docs.python.org/3/library/contextlib.html#contextlib.redirect_stdout), which works to the same effect. – sytech Jan 07 '21 at 08:33