[tmva][sofie] Don't expect warnings from PyTorch

tutorials/machine_learning/TMVA_SOFIE_Models.py

@@ -11,38 +11,14 @@
 ### \macro_output
 ### \author Lorenzo Moneta
 
-import contextlib
 import inspect
 import os
-import warnings
 
 import numpy as np
 import ROOT
 import torch
 import torch.nn as nn
 
-
-@contextlib.contextmanager
-def expect_warning(category, message):
-    # Silence a known third-party warning and raise if it stops firing.
-
-    # Notifies us to drop the workaround once the upstream library is fixed.
-    with warnings.catch_warnings(record=True) as caught:
-        warnings.simplefilter("always")
-        yield
-    seen = False
-    for w in caught:
-        if issubclass(w.category, category) and message in str(w.message):
-            seen = True
-        else:
-            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)
-    if not seen:
-        raise RuntimeError(
-            f"Expected {category.__name__} containing {message!r} was not "
-            "emitted. This tutorial's workaround can probably be removed."
-        )
-
-
 ## generate and train PyTorch models with different architectures
 
 
@@ -101,14 +77,7 @@ def filtered_kwargs(func, **candidate_kwargs):
     )
     print("calling torch.onnx.export with parameters", kwargs)
 
-    try:
-        # torch.onnx.export (dynamo path) pickles its export program through
-        # copyreg, which still references the deprecated LeafSpec. The warning
-        # is emitted from inside PyTorch and cannot be avoided from user code.
-        with expect_warning(FutureWarning, "isinstance(treespec, LeafSpec)"):
-            torch.onnx.export(model, dummy_x, modelFile, **kwargs)
-    except TypeError as e:
-        raise RuntimeError("Cannot export model from pytorch to ONNX - with version " + torch.__version__) from e
+    torch.onnx.export(model, dummy_x, modelFile, **kwargs)
 
     print("model exported to ONNX as", modelFile)
     return modelFile

tutorials/machine_learning/TMVA_SOFIE_ONNX.py

@@ -11,37 +11,14 @@
 ## \macro_output
 ## \author Lorenzo Moneta
 
-import contextlib
 import inspect
-import warnings
 
 import numpy as np
 import ROOT
 import torch
 import torch.nn as nn
 
 
-@contextlib.contextmanager
-def expect_warning(category, message):
-    # Silence a known third-party warning and raise if it stops firing.
-
-    # Notifies us to drop the workaround once the upstream library is fixed.
-    with warnings.catch_warnings(record=True) as caught:
-        warnings.simplefilter("always")
-        yield
-    seen = False
-    for w in caught:
-        if issubclass(w.category, category) and message in str(w.message):
-            seen = True
-        else:
-            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)
-    if not seen:
-        raise RuntimeError(
-            f"Expected {category.__name__} containing {message!r} was not "
-            "emitted. This tutorial's workaround can probably be removed."
-        )
-
-
 def CreateAndTrainModel(modelName):
 
     model = nn.Sequential(nn.Linear(32, 16), nn.ReLU(), nn.Linear(16, 8), nn.ReLU(), nn.Linear(8, 2), nn.Softmax(dim=1))
@@ -84,14 +61,7 @@ def filtered_kwargs(func, **candidate_kwargs):
     )
     print("calling torch.onnx.export with parameters", kwargs)
 
-    try:
-        # torch.onnx.export (dynamo path) pickles its export program through
-        # copyreg, which still references the deprecated LeafSpec. The warning
-        # is emitted from inside PyTorch and cannot be avoided from user code.
-        with expect_warning(FutureWarning, "isinstance(treespec, LeafSpec)"):
-            torch.onnx.export(model, dummy_x, modelFile, **kwargs)
-    except TypeError as e:
-        raise RuntimeError("Cannot export model from pytorch to ONNX - with version " + torch.__version__) from e
+    torch.onnx.export(model, dummy_x, modelFile, **kwargs)
 
     print("model exported to ONNX as", modelFile)
     return modelFile

tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py

@@ -12,37 +12,14 @@
 ### \macro_code
 ### \macro_output
 
-import contextlib
 import inspect
-import warnings
 
 import numpy as np
 import ROOT
 import torch
 import torch.nn as nn
 
 
-@contextlib.contextmanager
-def expect_warning(category, message):
-    # Silence a known third-party warning and raise if it stops firing.
-
-    # Notifies us to drop the workaround once the upstream library is fixed.
-    with warnings.catch_warnings(record=True) as caught:
-        warnings.simplefilter("always")
-        yield
-    seen = False
-    for w in caught:
-        if issubclass(w.category, category) and message in str(w.message):
-            seen = True
-        else:
-            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)
-    if not seen:
-        raise RuntimeError(
-            f"Expected {category.__name__} containing {message!r} was not "
-            "emitted. This tutorial's workaround can probably be removed."
-        )
-
-
 def PrepareData():
     # get the input data
     inputFile = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
@@ -135,14 +112,7 @@ def filtered_kwargs(func, **candidate_kwargs):
     )
     print("calling torch.onnx.export with parameters", kwargs)
 
-    try:
-        # torch.onnx.export (dynamo path) pickles its export program through
-        # copyreg, which still references the deprecated LeafSpec. The warning
-        # is emitted from inside PyTorch and cannot be avoided from user code.
-        with expect_warning(FutureWarning, "isinstance(treespec, LeafSpec)"):
-            torch.onnx.export(model, dummy_x, modelFile, **kwargs)
-    except TypeError as e:
-        raise RuntimeError("Cannot export model from pytorch to ONNX - with version " + torch.__version__) from e
+    torch.onnx.export(model, dummy_x, modelFile, **kwargs)
 
     print("model exported to ONNX as", modelFile)
     return modelFile