cls_head 计算特征时维度不匹配？

[lichen14]

PseCo/fsc147/4_1_train_roi_head.py

Line 277 in 0192ee5

cls_outs = cls_head(features, anchor_boxes, query_features)

我在跑这行代码的时候总是会遇到embeddings和prompts特征维度不匹配的问题，导致

PseCo/models.py

Line 25 in 0192ee5

pred_logits = (embeddings * prompts.unsqueeze(1)).sum(dim=-1)

这里无法直接矩阵乘

比如：

• 计算
  

  PseCo/fsc147/4_1_train_roi_head.py

  Line 277 in 0192ee5

  cls_outs = cls_head(features, anchor_boxes, query_features)

  
  时，embeddings 和 prompts.unsqueeze(1) 的shape 分别是 torch.Size([8192, 3, 512]) torch.Size([8192, 1, 3, 512]) ，这样直接计算会显存爆炸，需要384G ** 我试过把prompts.unsqueeze(1) 去掉unsqueeze(1) ，这样不报显存OOM的bug了，但是下面的代码会报错：
• 计算
  

  PseCo/fsc147/4_1_train_roi_head.py

  Line 283 in 0192ee5

  cls_outs2 = cls_head(features_a, clip_boxes, clip_target_features)

  
  时，embeddings 和 prompts的shape是torch.Size([1536, 3, 512]) torch.Size([1536, 512])，这时候就会有 RuntimeError: The size of tensor a (3) must match the size of tensor b (1536) at non-singleton dimension 1 。此时如果加上prompts.unsqueeze(1)就不会报错了

所以，是否是代码历史版本的问题？有无最新版的可行代码呀？