Add torchinfo package to summarize a pytorch model (#1220)

http://b/257562539

Author: Philmod

Dockerfile.tmpl

@@ -559,7 +559,8 @@ RUN pip install pytorch-ignite \
         tables \
         openpyxl \
         timm \
-        pycolmap && \
+        pycolmap \
+        torchinfo && \
     /tmp/clean-layer.sh
 
 # Download base easyocr models.

tests/test_torchinfo.py

@@ -0,0 +1,11 @@
+import unittest
+
+import torch.nn as tnn
+
+from torchinfo import summary
+
+class TestTorchinfo(unittest.TestCase):
+    def test_info(self):
+        model = tnn.Linear(5,3)
+        s = summary(model)
+        self.assertEqual(1, len(s.summary_list))