a3d053ee6de9b98d205cf3256e2d7f70c326d88a,gluon/gluoncv2/models/superpointnet.py,SPDescriptor,hybrid_forward,#SPDescriptor#Any#Any#Any#,249

Before Change

        desc_map = desc_map.reshape(shape=(0, 0, -1))
        desc_map = desc_map.transpose(axes=(0, 2, 1))

        desc_map_list = []

        def slice_desc_map(data, _):
            desc_map_list.append(data)
            return data, []
        F.contrib.foreach(slice_desc_map, desc_map, [])

        pts_ravel_list = []

        def ravel_pts(data, _):
            pts_ravel_list.append(F.ravel_multi_index(data, shape=(self.in_size[0], self.in_size[1])))
            return data, []
        pts_tr = pts.transpose(axes=(0, 2, 1))
        F.contrib.foreach(ravel_pts, pts_tr, [])

        desc_map_sorted_list = []
        for desc_map_i, pts_coord_ravel_i in zip(desc_map_list, pts_ravel_list):
            desc_map_sorted_i = F.take(desc_map_i, pts_coord_ravel_i)
            desc_map_sorted_list.append(desc_map_sorted_i)

        desc_map_sorted = F.stack(*desc_map_sorted_list)

        return desc_map_sorted

After Change

        pts_tr = pts.transpose(axes=(0, 2, 1))
        for i in range(self.batch_size):
            desc_map_i = desc_map[i]
            pts_tr_i = pts_tr[i].reshape(shape=(2, -1))
            pts_ravel_i = F.ravel_multi_index(pts_tr_i, shape=(self.in_size[0], self.in_size[1]))
            desc_map_sorted_i = F.take(desc_map_i, pts_ravel_i)
            desc_map_sorted_list.append(desc_map_sorted_i)
        desc_map_sorted = F.stack(*desc_map_sorted_list)

Instances