add a tool for combining visualization

Summary: Pull Request resolved: #169 add a tool for combining visualization Reviewed By: zechenghe Differential Revision: D46949405 fbshipit-source-id: 62bdbebafd37e9526a3690ec90bde0510af08a56
facebookresearch · Jun 29, 2023 · f17e4fc · f17e4fc
1 parent 1ed6402
commit f17e4fc
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/mobile_cv/torch/tests/utils_pytorch/test_state_dict_name_matcher.py b/mobile_cv/torch/tests/utils_pytorch/test_state_dict_name_matcher.py
@@ -54,3 +54,31 @@ def test_matcher_shape(self):
                 "model.m1.bb.weight": "model.m1.dda.weight",
             },
         )
+
+    def test_list_mapping(self):
+        list1 = [
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_06-16-2023_20:12:49.png",
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_hint_06-16-2023_20:12:51.png",
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_mask_06-16-2023_20:12:53.png",
+            "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:10:05_0.763774618976614_06-16-2023_20:10:07.png",
+        ]
+
+        list2 = [
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_06-16-2023_20:14:06.png",
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_hint_06-16-2023_20:14:08.png",
+            "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_mask_06-16-2023_20:14:10.png",
+            "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:11:22_0.763774618976614_06-16-2023_20:11:24.png",
+            "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:11:22_0.763774618976614_hint_06-16-2023_20:11:28.png",
+            "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:11:22_0.763774618976614_mask_06-16-2023_20:11:30.png",
+        ]
+
+        mapping = sdnm.get_list_mapping(list1, list2)
+        self.assertDictEqual(
+            mapping,
+            {
+                "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_06-16-2023_20:12:49.png": "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_06-16-2023_20:14:06.png",
+                "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_hint_06-16-2023_20:12:51.png": "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_hint_06-16-2023_20:14:08.png",
+                "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:12:48_0.8357651039198697_mask_06-16-2023_20:12:53.png": "gpu0_l_ggmte_zp_ae_g_ziemnv_iv_g_irmcw_kged_abcdefghij_006-16-2023_20:14:05_0.8357651039198697_mask_06-16-2023_20:14:10.png",
+                "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:10:05_0.763774618976614_06-16-2023_20:10:07.png": "gpu0_b_voree_vbtmt_otym_t_trmrtr_gmeritm_wx_n_rmgri_006-16-2023_20:11:22_0.763774618976614_06-16-2023_20:11:24.png",
+            },
+        )
diff --git a/mobile_cv/torch/utils_pytorch/state_dict_name_matcher.py b/mobile_cv/torch/utils_pytorch/state_dict_name_matcher.py
@@ -217,3 +217,29 @@ def get_match_keys_bipartite(fp32_items: List[Item], qat_items: List[Item]):
         # print(f"Best matching cost: {cost_sum}")
 
     return ret
+
+
+def get_list_mapping(list1: List[str], list2: List[str]) -> Dict[str, str]:
+    """Get the best mapping between 2 lists of strings"""
+    ret = []
+
+    dim1 = len(list1)
+    dim2 = len(list2)
+
+    cost = np.zeros((dim1, dim2))
+    for ii in range(dim1):
+        for jj in range(dim2):
+            cost[ii, jj] = levenshtein(list1[ii], list2[jj])
+
+    from scipy.optimize import linear_sum_assignment  # @manual
+
+    # bipartite matching with minimal cost
+    row_ind, col_ind = linear_sum_assignment(cost)
+    ret = {list1[ri]: list2[ci] for ri, ci in zip(row_ind, col_ind)}
+    cost_sum = cost[row_ind, col_ind].sum()
+    if cost_sum > 0:
+        for list1_key, list2_key in ret.items():
+            if list1_key != list2_key:
+                print(f"{list1_key} -> {list2_key}")
+
+    return ret