Install onnx package (Kaggle#1170)

http:https://b/235510037
RodrigoGonzalez · Jun 15, 2022 · aef6592 · aef6592
1 parent 5a4cb0e
commit aef6592
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
@@ -511,6 +511,7 @@ RUN pip install flashtext && \
  # ipympl adds interactive widget support for matplotlib
  pip install ipympl==0.7.0 && \
  pip install pandarallel && \
+ pip install onnx && \
  /tmp/clean-layer.sh
 
 # Download base easyocr models.

diff --git a/tests/data/mnist-8.onnx b/tests/data/mnist-8.onnx
diff --git a/tests/test_onnx.py b/tests/test_onnx.py
@@ -0,0 +1,8 @@
+import unittest
+
+import onnx
+
+class TestOnnx(unittest.TestCase):
+ def test_load(self):
+ model = onnx.load("/input/tests/data/mnist-8.onnx")
+ self.assertIn("CNTKGraph", onnx.helper.printable_graph(model.graph))