Shape mismatch error when increase num_frames

Hi, thanks again for your help to solve my previous questions!

While training the model in stage 3, I tried to increase the num_frames in the config file, but the program ran into error when processing certain video.

I used some of the IT datasets as you provided on HuggingFace, such as ss_v2, videochat1, textvr, etc. No other datasets were included. Here is my configuration (the others were kept the same):

num_frames = 23
batch_size = 1
img_size = 224

The error information is as follows:

/workspace/data/datasets/videochat2/ss_v2/20bn-something-something-v2/15910.webm
ITERATION VIDEO_SHAPE:  torch.Size([1, 20, 3, 224, 224])
Model forward VIDEO_SHAPE:  torch.Size([1, 20, 3, 224, 224])
When loading data, NUM_FRAMES:  23 , VLEN 45
Dataset __getitem__ right before return, VIDEO SHAPE:  torch.Size([23, 3, 224, 224])
PATCH EMBED SHAPE:  torch.Size([1, 3920, 1024])
POS EMBED SHAPE:  torch.Size([1, 4508, 1024])
Traceback (most recent call last):
  File "/workspace/data/code/vc2_f20_i224/train_it.py", line 219, in <module>
    main(cfg)
  File "/workspace/data/code/vc2_f20_i224/train_it.py", line 167, in main
    global_step = train(
  File "/workspace/data/code/vc2_f20_i224/train_it.py", line 61, in train
    loss_dict = model(image, text, instruction)
  File "/opt/conda/envs/vc2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/workspace/data/code/vc2_f20_i224/models/videochat2_it.py", line 233, in forward
    img_embeds, use_image = self.encode_img(image, instruction)
  File "/workspace/data/code/vc2_f20_i224/models/videochat2_it.py", line 186, in encode_img
    image_embeds = self.vision_encoder(image, use_image)
  File "/opt/conda/envs/vc2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/workspace/data/code/vc2_f20_i224/models/blip2/vit.py", line 405, in forward
    x_vis = self.encoder(x, use_image) # [B, N_vis, C_e]
  File "/opt/conda/envs/vc2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/workspace/data/code/vc2_f20_i224/models/blip2/vit.py", line 327, in forward
    x_vis = self.forward_features(x, use_image)
  File "/workspace/data/code/vc2_f20_i224/models/blip2/vit.py", line 311, in forward_features
    x = x + self.pos_embed.type_as(x).to(x.device).clone().detach()
RuntimeError: The size of tensor a (3920) must match the size of tensor b (4508) at non-singleton dimension 1

In order to figure out why the shape mismatch error occurs, I manually add some print functions to check the tensor shape at each step. The output of each step is shown above (after "Traceback").

1. In "train_it.py", def train, I printed the video shape right after enumerating the iterator.

for i, (media_type, (image, text, instruction, _)) in enumerate(iterator):
        print("ITERATION VIDEO_SHAPE: ", image.shape)
        image = image.to(device, non_blocking=True)

...

2. In "videochat2_it.py", def forward, I printed the video shape right after receiving the video.

def forward(self, image, text_input, instruction):
        print("Model forward VIDEO_SHAPE: ", image.shape)
        img_embeds, use_image = self.encode_img(image, instruction)
        batch_size, img_len, _ = img_embeds.shape

...
3. In "it_dataset.py", class ITVidTrainDataset, def __getitem__, I printed the video shape right before return the video.

def __getitem__(self, index):
        try:
            ann = self.get_anno(index)
            print(ann["image"])
            msg = ""
            clip = None
            if "start" in ann and "end" in ann:
                clip = [ann["start"], ann["end"]]
            video, index, sec = self.load_and_transform_media_data_video(index, ann["image"], return_fps=True, clip=clip)
            if self.add_second_msg:
                # " " should be added in the start and end
                msg = f" The video contains {len(sec)} frames sampled at {', '.join(sec)} seconds. "
            conversation, instruction = self.process_qa(ann["qa"], msg)
            print("Dataset __getitem__ right before return, VIDEO SHAPE: ", video.shape)
            return video, conversation, instruction, index
        except Exception as e:
            logger.warning(f"Caught exception {e} when loading video {ann['image']}")
            index = np.random.randint(0, len(self))
            return self.__getitem__(index)

...

4. In "video_utils.py", def get_frame_indices, I printed the number of frames and the length of video after loading.

def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
    if sample in ["rand", "middle"]: # uniform sampling
        print("When loading data, NUM_FRAMES: ", num_frames, ", VLEN", vlen)
        acc_samples = min(num_frames, vlen)
        # split the video into `acc_samples` intervals, and sample from each interval.
        intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
        ranges = []
        for idx, interv in enumerate(intervals[:-1]):
            ranges.append((interv, intervals[idx + 1] - 1))

...

5. In "vit.py", class PretrainVisionTransformerEncoder, def forward_features, I printed the video shape after patch embed and the corresponding pos embed shape.

def forward_features(self, x, use_image=False):
        x = self.patch_embed(x)
        print("PATCH EMBED SHAPE: ", x.shape)
        
        if use_image:
            x = x + self.img_pos_embed.type_as(x).to(x.device).clone().detach()
        else:
            print("POS EMBED SHAPE: ", self.pos_embed.shape)
            x = x + self.pos_embed.type_as(x).to(x.device).clone().detach()

...

The results may indicate the video shape changes after the data loader. Also, the error happens only on certain videos. I tried to remove the video that causes an error, then the error happens for another video.

The most confused thing is that most videos are OK but there is a random video that causes error still waiting for me :(

BTW, when I set the num_frames=16, there is no error. I tried 18, 20, 24, there are also errors.

Is there any solution for this issue? Many thanks for your help

Reproduce stage 3 training with num_frames=23 using the listed ss_v2 video, then trace the tensor shapes through train_it.py, it_dataset.py, video_utils.py, and models/blip2/vit.py. Compare the sampled frame count, patch embedding, and positional embedding for affected videos and verify that training completes for the supported num_frames settings without a shape mismatch.