TestWriteONNX Class Reference

Inheritance diagram for TestWriteONNX:

Collaboration diagram for TestWriteONNX:

Public Member Functions
	create_and_save (self, n_outputs, filename, weights)
	test_singleclass (self)
	test_multiclass (self)

Detailed Description

Tests for writing ONNX MVA weightfiles. In addition to testing the writing
mechanism, these serve the purpose of creating test files for other unit tests.

Definition at line 37 of file test_write_onnx.py.

Member Function Documentation

create_and_save()

create_and_save	(		self,
			n_outputs,
			filename,
			weights )

Setup some example options for an ONNX MVA weightfile, save it and compare to reference

Definition at line 42 of file test_write_onnx.py.

    def create_and_save(self, n_outputs, filename, weights):
        """
        Setup some example options for an ONNX MVA weightfile, save it and compare to reference
        """
 
        # like in mva/tests/all_classifiers.py
        variables = [
            "p",
            "pz",
            "daughter(0, p)",
            "daughter(0, pz)",
            "daughter(1, p)",
            "daughter(1, pz)",
            "chiProb",
            "dr",
            "dz",
            "daughter(0, dr)",
            "daughter(1, dr)",
            "daughter(0, chiProb)",
            "daughter(1, chiProb)",
            "daughter(0, kaonID)",
            "daughter(0, pionID)",
            "daughterAngle(0, 1)",
        ]
        model = nn.Linear(len(variables), n_outputs)
        model.load_state_dict(weights)
        general_options = basf2_mva.GeneralOptions()
        general_options.m_datafiles = basf2_mva.vector("dummy")
        general_options.m_identifier = "Simple"
        general_options.m_treename = "tree"
        general_options.m_variables = basf2_mva.vector(*variables)
        specific_options = basf2_mva.ONNXOptions()
        with b2test_utils.clean_working_directory():
            save_onnx(model, general_options, specific_options, filename)
            with open(filename) as f:
                xml_new = f.read()
        ref_path = Path(basf2.find_file("mva/methods/tests")) / filename
        try:
            with open(ref_path) as f:
                xml_ref = f.read()
        except FileNotFoundError:
            # if the file does not exist, recreate it, but still fail the test
            with open(ref_path, "w") as f:
                f.write(xml_new)
            raise Exception(f"Wrote new reference file {str(ref_path)}")
        self.assertEqual(xml_new, xml_ref)
 

test_multiclass()

test_multiclass

(

self

)

Write example for multiclass outputs

Definition at line 105 of file test_write_onnx.py.

    def test_multiclass(self):
        """
        Write example for multiclass outputs
        """
        self.create_and_save(
            2,
            "ONNX_multiclass.xml",
            {
                "weight": torch.tensor(
                    [[0.1911,  0.2075, -0.0586,  0.2297, -0.0548,  0.0504, -0.1217,  0.1468,
                      0.2204, -0.1834,  0.2173,  0.0468,  0.1847,  0.0339,  0.1205, -0.0353],
                     [0.1927,  0.0370, -0.1167,  0.0637, -0.1152, -0.0293, -0.1015,  0.1658,
                      -0.1973, -0.1153, -0.0706, -0.1503,  0.0236, -0.2469,  0.2258, -0.2124]]
                ),
                "bias": torch.tensor([0.1930, 0.0416])
            }
        )
 
 

test_singleclass()

test_singleclass

(

self

)

Write example for single output

Definition at line 89 of file test_write_onnx.py.

    def test_singleclass(self):
        """
        Write example for single output
        """
        self.create_and_save(
            1,
            "ONNX.xml",
            {
                "weight": torch.tensor(
                    [[0.1911,  0.2075, -0.0586,  0.2297, -0.0548,  0.0504, -0.1217,  0.1468,
                      0.2204, -0.1834,  0.2173,  0.0468,  0.1847,  0.0339,  0.1205, -0.0353]]
                ),
                "bias": torch.tensor([0.1927])
            }
        )
 

---

The documentation for this class was generated from the following file:

• mva/methods/tests/test_write_onnx.py