Rewrite image embedding to remove the in-place op

image_embedding_phi3_v.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from datetime import datetime
-
 import torch
 from torch import nn
 from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig
@@ -28,6 +26,9 @@ except ImportError:
 
 logger = logging.get_logger(__name__)
 
+
+MAX_INPUT_ID = int(1e9)
+
 CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(
   attention_dropout=0.0,
   dropout=0.0,
@@ -179,32 +180,44 @@ class Phi3ImageEmbedding(nn.Module):
             patch_feature = img_feature[:, 1:]
             return patch_feature
 
-        if TYPE_FEATURE == "cls_patch":
-            return img_feature
-
         raise NotImplementedError
 
-    def forward(self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None) -> torch.FloatTensor:
-
-        MAX_INPUT_ID = int(1e9)
-        img_embeds = pixel_values
-        img_sizes = image_sizes
-
-        if self.img_features is not None:
-            img_embeds = self.img_features.clone()
-            self.img_features = None
-
-        if self.img_sizes is not None:
-            img_sizes = self.img_sizes
-
+    def forward(
+        self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None
+    ) -> torch.FloatTensor:
         input_shape = input_ids.size()
         input_ids = input_ids.view(-1, input_shape[-1])
 
-        with torch.no_grad():
-            positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=False)
+        # positions for image tokens
+        positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)
+        has_image = len(positions[0].tolist()) > 0
+        input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()
+        hidden_states = self.wte(input_ids)
+
+        if has_image:
+            assert self.use_hd_transform
+            num_images, num_crops, c, h, w = pixel_values.shape
+            assert c == 3 and h == w == 336
+            img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(
+                num_images, num_crops, -1, self.image_dim_out
+            )
+            image_features_proj = self.hd_feature_transform(img_features, image_sizes)
+            hidden_states = hidden_states.index_put(
+                positions, image_features_proj, accumulate=False
+            )
+
+        if self.drop is not None:
+            hidden_states = self.drop(hidden_states)
 
-        select = False
+        return hidden_states
 
+    def hd_feature_transform(self, image_features, image_sizes):
+        """
+        image_features: (num_images, num_crops+1, 24*24, 1024)
+        """
+        assert (
+            self.hd_transform_order == 'sub_glb'
+        ), f'hd_transform_order `{self.hd_transform_order}` not implemented'
         if isinstance(self.img_projection, nn.Sequential):
             target_device = self.img_projection[0].bias.device
             target_dtype = self.img_projection[0].bias.dtype
@@ -212,135 +225,98 @@ class Phi3ImageEmbedding(nn.Module):
             target_device = self.img_projection.bias.device
             target_dtype = self.img_projection.bias.dtype
 
-        if len(positions.tolist()) > 0:
-            with torch.no_grad():
-                g_values = abs(input_ids[positions[:, 0], positions[:, 1]])
-
-            if self.use_hd_transform and img_sizes is not None and len(img_sizes):
-                hd_transform = True
-                assert img_embeds.ndim == 5, f'img_embeds size: {img_embeds.size()}, expect 5D tensor for hd transform'
-                # img_embeds: (num_images, max_num_crops, 3, H, W)
-                # img_sizes: (num_images, 2).view(1, -1)
-
-                start_time = datetime.now()
-                bs = img_embeds.shape[0]
-                # Nx(HW)xC
-                img_features = self.get_img_features(img_embeds.flatten(0, 1))
-                base_feat_height = base_feat_width = int(img_features.shape[1] ** 0.5)
-
-                assert base_feat_height == 24 and base_feat_width == 24, f'base_feat_height: {base_feat_height}, base_feat_width: {base_feat_width}, expect 24x24 features for hd transform'
-
-                # bs x max_num_crops x (24x24) x C
-                img_features = img_features.view(bs, -1, base_feat_height * base_feat_width, self.image_dim_out)
-                C = self.image_dim_out
-                H = base_feat_height
-
-                output_imgs = []
-                output_len = []
-                # training is tensor, inference is list
-                if isinstance(img_sizes, torch.Tensor):
-                    img_sizes = img_sizes.view(-1, 2)
-                for _bs in range(bs):
-                    h, w = img_sizes[_bs]
-                    h = h // 336
-                    w = w // 336
-                    B_ = h * w
-
-                    # 1 x (24x24) x 1024
-                    global_img_feature = img_features[_bs, :1]
-
-                    # 1 x 12 x 12 x 4096
-                    glb_img = global_img_feature.reshape(1,H,H,C).reshape(1,H//2,2,H//2,2,C).contiguous().permute(0,1,3,2,4,5).reshape(1,H//2,H//2,4*C).contiguous()
-                    temp_glb_GN = self.sub_GN.repeat(1, H//2, 1, 1)
-
-                    # 1 x 156 x 4096
-                    glb_img = torch.cat([glb_img, temp_glb_GN], dim=2).reshape(1,-1,4*C)
-
-                    # (max_num_crops-1) x (12x12) x C
-                    sub_img = img_features[_bs, 1:]
-                    # 16x574x1024
-                    # get rid of padding sub_img
-                    sub_img = sub_img[:B_]
-
-                    # (num_crops, 12, 2, 12, 2, 1024) -> (num_crops, 12, 12, 2, 2, 1024) -> (num_crops, 12*12, 4*1024)
-                    sub_img = sub_img.reshape(B_,H,H,C).reshape(B_,H//2,2,H//2,2,C).contiguous().permute(0,1,3,2,4,5).reshape(B_,-1,4*C).contiguous()
-                    sub_img = sub_img.reshape(1, h, w, 12, 12, -1).permute(0,1,3,2,4,5).reshape(1,h*12,w*12,4*C)
-                    temp_sub_GN = self.sub_GN.repeat(1, h*12, 1, 1)
-                    sub_img = torch.cat([sub_img, temp_sub_GN], dim=2).reshape(1,-1,4*C)
-                    # (1, num_img_tokens, 1024*4)
-
-                    # glb + sub
-                    if self.hd_transform_order == 'glb_sub':
-                        output_imgs.append(torch.cat([glb_img, self.glb_GN, sub_img], dim=1))
-                    elif self.hd_transform_order == 'sub_glb':
-                        output_imgs.append(torch.cat([sub_img, self.glb_GN, glb_img], dim=1))
-                    else:
-                        raise NotImplementedError(f'hd_transform_order = {self.hd_transform_order}, not implemented')
-
-                    temp_len = int((h*w+1)*144 + 1 + (h+1)*12)
-                    assert temp_len == output_imgs[-1].shape[1], f'temp_len: {temp_len}, output_imgs[-1].shape[1]: {output_imgs[-1].shape[1]}'
-                    output_len.append(temp_len)
-
-                num_img_tokens = output_len
-                img_set_tensor = []
-                for _output_img in output_imgs:
-                    img_feature_proj = self.img_projection(_output_img.to(target_device).to(target_dtype))
-                    img_set_tensor.append(img_feature_proj)
-                logger.info(f'img_embeds size: {img_embeds.size()}, image sizes: {img_sizes} loading time {datetime.now() - start_time}')
-            elif img_embeds.ndim == 4:
-                selected_g_values = g_values[::self.num_img_tokens]
-                assert len(img_embeds) == len(selected_g_values), f'img_embeds size: {img_embeds.size()}, selected_g_values size: {len(selected_g_values)}, selected_g_value {selected_g_values}'
-                start_time = datetime.now()
-                tt = (
-                    self.get_img_features(img_embeds)
-                    .to(target_device)
-                    .to(target_dtype)
-                    .reshape(-1, self.image_dim_out)
-                )
-                logger.info(f'img_embeds size: {img_embeds.size()}, loading time {datetime.now() - start_time}')
-                img_set_tensor = self.img_projection(tt)  # adapted visual features.
-            elif img_embeds.ndim == 3:
-                selected_g_values = g_values[::self.num_img_tokens]
-                assert len(img_embeds) == len(selected_g_values), f'img_embeds size: {img_embeds.size()}, selected_g_values size: {len(selected_g_values)}, selected_g_value {selected_g_values}'
-                tt = (
-                    img_embeds
-                    .to(target_device)
-                    .to(target_dtype)
-                    .view(-1, self.image_dim_out)
-                )
-                img_set_tensor = self.img_projection(tt)  # adapted visual features.
-            else:
-                raise NotImplementedError
-            select = True
-
-        with torch.no_grad():
-            input_ids.clamp_min_(0).clamp_max_(self.vocab_size)
-
-        hidden_states = self.wte(input_ids)
-
-        if select:
-            if hd_transform:
-                idx = 0
-                for i, cnt in enumerate(num_img_tokens):
-                    hidden_states[positions[idx, 0], positions[idx, 1] : positions[idx, 1] + cnt] = (
-                        img_set_tensor[i]
-                        .to(hidden_states.dtype)
-                        .to(hidden_states.device)
-                        )
-                    idx += cnt
-            else:
-                idx = 0
-                assert len(selected_g_values) * self.num_img_tokens == len(img_set_tensor), f'len(selected_g_values) * self.num_img_tokens = {len(selected_g_values) * self.num_img_tokens}, len(img_set_tensor) = {len(img_set_tensor)}'
-                for i, g in enumerate(selected_g_values):
-                    cnt = self.num_img_tokens
-                    hidden_states[positions[idx, 0], positions[idx, 1] : positions[idx, 1] + cnt] = (
-                        img_set_tensor[i * cnt : (i + 1) * cnt]
-                        .to(hidden_states.dtype)
-                        .to(hidden_states.device)
-                        )
-                    idx += cnt
-
-        if self.drop is not None:
-            hidden_states = self.drop(hidden_states)
-
-        return hidden_states
+        global_image_features = image_features[:, 0]  # (num_images, 24*24, 1024)
+        # global feature can be viewed as a special HD case with num_crops 1x1
+        global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
+        global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
+
+        all_image_embeddings = []
+        # need a for loop to process each image because of different image sizes
+        # (patch arrangement is different for each image)
+        for i, img_size in enumerate(image_sizes):
+            h, w = img_size
+            h_crop = h // 336
+            w_crop = w // 336
+            num_crops = h_crop * w_crop
+
+            # NOTE: real num_crops is padded
+            # (num_crops, 24*24, 1024)
+            sub_image_features = image_features[i, 1 : 1 + num_crops]
+            sub_image_features_hd = self.reshape_hd_patches_2x2merge(
+                sub_image_features, h_crop, w_crop
+            )
+            sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
+
+            # [sub features, separator, global features]
+            all_image_embeddings.extend(
+                [
+                    sub_image_features_hd_newline.squeeze(0),  # (h_crop*12*(w_crop*12+1), 4096)
+                    self.glb_GN.squeeze(0),
+                    global_image_features_hd_newline[i],
+                ]
+            )
+
+        image_features_proj = self.img_projection(
+            torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)
+        )
+
+        return image_features_proj
+
+    def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
+        """
+        image_features: (num_images*num_crops, 24*24, 1024)
+        output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops
+        """
+        N, L, C = image_features.shape
+        assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0
+        num_images = N // (h_crop * w_crop)
+        H = int(L**0.5)
+        image_features_hd = (
+            image_features.reshape(N, H, H, C)  # N, 24, 24, 1024
+            .reshape(N, H // 2, 2, H // 2, 2, C)  # N, 12, 2, 12, 2, 1024
+            .permute(0, 1, 3, 2, 4, 5)  # N, 12, 12, 2, 2, 1024
+            .reshape(N, -1, 4 * C)  # N, 144, 4096
+            .reshape(
+                num_images, h_crop, w_crop, H // 2, H // 2, -1
+            )  # n_img, h_crop, w_crop, 12, 12, 4096
+            .permute(0, 1, 3, 2, 4, 5)  # n_img, h_crop, 12, w_crop, 12, 4096
+            .reshape(
+                num_images, h_crop * H // 2, w_crop * H // 2, 4 * C
+            )  # n_img, h_crop*12, w_crop*12, 4096
+        )
+
+        # alternative implementation using einops
+        # from einops import rearrange
+        # image_features_nhwc = rearrange(
+        #     image_features,
+        #     'N (H W) c -> N H W c',
+        #     H=H,
+        #     W=H,
+        # )
+        # image_features_2x2merge = rearrange(
+        #     image_features_nhwc,
+        #     'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',
+        #     h_pool=2,
+        #     w_pool=2,
+        # )
+        # image_features_hd = rearrange(
+        #     image_features_2x2merge,
+        #     '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',
+        #     h_crop=h_crop,
+        #     w_crop=w_crop,
+        # )
+
+        return image_features_hd
+
+    def add_image_newline(self, image_features_hd):
+        """
+        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
+        output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
+        """
+        num_images, h, w, hid_dim = image_features_hd.shape
+        # add the newline token to the HD image feature patches
+        newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1)  # (n_img, h, 1, hid_dim)
+        image_features_hd_newline = torch.cat(
+            [image_features_hd, newline_embeddings], dim=2
+        ).reshape(num_images, -1, hid_dim)
+        return image_features_hd_newline