warnings.warn is too spammy in TorchScript

[dzhulgakov]

🐛 Bug

TorchScript supports warnings.warn syntax, however it doesn't respect the Python defaults (which is print the warning only once per invocation point as pointed out by stack_level): https://docs.python.org/3/library/warnings.html#the-warnings-filter

It can cause a lot of spam, especially for deployed inference models.

To Reproduce

Sample notebook: https://colab.research.google.com/drive/1Zcm4F3rk58Q8aEBtfxSgHOKGJ9Ib2Uhs?usp=sharing

import torch

def foo(x):
  for i in range(10):
    x = torch.nn.functional.softmax(x)  # generates a warning inside

foo(torch.rand(5))

# warning is printed once

torch.jit.script(foo)(torch.rand(5))

# warning is printed 10 times

Expected behavior

Ideally, we'd emulate Python's default syntax directly, i.e. do the dedup based on the stack trace - the warn instruction in TorchScript actually does receive stack_level but ignores it.

As a simple fix though we can pretend that log_level is 0 - i.e. log only the first invocation of the unique instruction (e.g. can be done by adding a mutable flag to Code instance in Interpreter for each warn instruction.

Note, that TORCH_WARN_ONCE in C++ is not a good fit here - it'd put all warning to the same line in C++ which is not desired.

cc @ezyang @gchanan @zou3519 @gmagogsfm

[gmagogsfm]

Thanks for reporting this issue. I think we can improve this bad user experience but it is tricky to match exact Python behavior.

As you mentioned, Python relies on a callstack and stack_level to identify unique instances of warnings.warn. In TorchScript, we don't have the same callstack as Python's since calls may be inlined.

However, we could alternatively make sure that "each warning statement is triggered only once" rather than "each warning statement, identified by its stack_level, is triggered only once". Would that be useful enough? @dzhulgakov

Specifically, when emitting aten::warn, we could attach a unique identifier to each instruction as an attribute. Interpreter then tracks whether a aten::warn has triggered before according to the unique identifier. @SplitInfinity @eellison

[dzhulgakov]

Yes, that's what I figured too and suggested in the description: we can pretend that stack_level is always 0 (or is it 1?) and just log once per the instruction (e.g. by keeping a flag in Code object for each warn instruction)

[eellison]

SGTM. We inline everything so every unique stack trace will have its own warning node. I briefly marked this as. high priority because this both annoying to the user and can measurably slow down models (see: pytorch/vision#2314)

[gmagogsfm]

Thanks for the confirmation. I will take action to implement this.