check the model info

yux-lab · Mar 2, 2025 · 4ac582a · 4ac582a
1 parent c96313c
commit 4ac582a
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/pix2tex/model/checkpoints/model_info.py b/pix2tex/model/checkpoints/model_info.py
@@ -0,0 +1,13 @@
+import torch
+
+pth_file_path = "weights.pth"
+
+# load
+state_dict = torch.load(pth_file_path, map_location=torch.device('cpu'))
+
+# print
+# print("Model parameters:")
+# for key, value in state_dict.items():
+#     print(f"{key}: {value.shape}")
+total_params = sum(value.numel() for value in state_dict.values())
+print(f"Total number of parameters: {total_params}")