Search code examples
pythonpytorchverbosity

Is there a way to force any function to not be verbose in Python?


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?


Solution

  • 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.