pfnet / pytorch-pfn-extras

Supplementary components to accelerate research and development in PyTorch
https://medium.com/pytorch/migration-from-chainer-to-pytorch-8ed92c12c8
MIT License
271 stars 52 forks source link

onnx's CVE-2024-27318 fix breaks `load_model` for stripped models #809

Closed risicle closed 6 months ago

risicle commented 6 months ago

As of https://github.com/onnx/onnx/pull/5917 (not present in a released version yet), load_model now seems to raise onnx.checker.ValidationError for missing files instead of OSError. Which breaks your trick used at https://github.com/pfnet/pytorch-pfn-extras/blob/f6b127063ec910b71788db2ae6ef96a3d89832b1/pytorch_pfn_extras/onnx/load.py#L29 to retry loading in "stripped" mode. And it's not just a matter of adding that exception to the handler, because it doesn't look like their ValidationError is helpful enough to provide an equivalent of e.filename.

This causes test_stripped_onnx_load_model to fail:

________________________ test_stripped_onnx_load_model _________________________

    @pytest.mark.filterwarnings("ignore:.*ONNX contains stripped .*:UserWarning")
    def test_stripped_onnx_load_model():
        model = Net()
        outdir = "out/stripped_load_model_test"
        tou.export_testcase(model, torch.rand(1, 1, 28, 28), outdir,
                            strip_large_tensor_data=True, training=True,
                            do_constant_folding=False)
>       tou.load_model(os.path.join(outdir, "model.onnx"))

tests/pytorch_pfn_extras_tests/onnx_tests/test_load_model.py:26: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
pytorch_pfn_extras/onnx/load.py:26: in load_model
    return onnx.load_model(f, format=format, load_external_data=load_external_data)
/...onnx.../lib/python3.11/site-packages/onnx/__init__.py:176: in load_model
    load_external_data_for_model(model, base_dir)
/...onnx.../lib/python3.11/site-packages/onnx/external_data_helper.py:66: in load_external_data_for_model
    load_external_data_for_tensor(tensor, base_dir)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

tensor = dims: 20
dims: 1
dims: 5
dims: 5
data_type: 1
name: "conv1.weight"
external_data {
  key: "location"
  value: "{\"type...6728, \"variance\": 0.013710332103073597}"
}
external_data {
  key: "length"
  value: "2000"
}
data_location: EXTERNAL

base_dir = '/build/source/out/stripped_load_model_test'

    def load_external_data_for_tensor(tensor: TensorProto, base_dir: str) -> None:
        """
        Loads data from an external file for tensor.
        Ideally TensorProto should not hold any raw data but if it does it will be ignored.

        Arguments:
            tensor: a TensorProto object.
            base_dir: directory that contains the external data.
        """
        info = ExternalDataInfo(tensor)
>       external_data_file_path = c_checker._resolve_external_data_location(  # type: ignore[attr-defined]
            base_dir, info.location, tensor.name
        )
E       onnx.onnx_cpp2py_export.checker.ValidationError: Data of TensorProto ( tensor name: conv1.weight) should be stored in /build/source/out/stripped_load_model_test/{"type": "stripped", "average": 0.005948468577116728, "variance": 0.013710332103073597}, but it doesn't exist or is not accessible.

/...onnx.../lib/python3.11/site-packages/onnx/external_data_helper.py:43: ValidationError
----------------------------- Captured stdout call -----------------------------
================ Diagnostic Run torch.onnx.export version 2.0.1 ================
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

So you might want to do something about that because everyone's going to want to upgrade to a version with this change as soon as there's a release what with it being the CVE-2024-27318 fix.

take-cheeze commented 6 months ago

@risicle Thank you for reporting! We'll fix this after new onnx package would be released!