预训练FaceNet模型输出维度为[2,512]引发torch.stack报错求助

核心问题原因 #

你的错误源于MTCNN(keep_all=True)会检测并保留图片中所有人脸：如果某帧视频里有2个人脸，输出的特征张量维度是[2,512]；只有1个人脸的帧输出是[1,512]。而torch.stack要求所有输入张量的维度完全一致，因此触发报错。

另外代码存在两处语法错误：

• resultlist.append(ImgFeatures(image_path) 缺少右括号
• troch.stack 拼写错误，正确写法是torch.stack

解决方案 #

根据需求，提供两种处理方式：

方式1：只保留每张图的第一张人脸 #

修改MTCNN参数为keep_all=False，每张图仅返回检测到的第一张人脸特征，所有张量维度统一为[512]，后续可正常使用torch.stack：

import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
from PIL import Image  # 需导入Image模块

resultlist = []
# 模型初始化放在函数外，避免重复创建提升效率
mtcnn = MTCNN(keep_all=False).eval()
model = InceptionResnetV1(pretrained='vggface2').eval()

def ImgFeatures(image_path):
    image = Image.open(image_path)
    cropped = mtcnn(image)
    if cropped is not None:  # 处理未检测到人脸的情况
        result = model(cropped).detach()
        return result
    else:
        # 可返回全0张量或跳过该帧，按需调整
        return torch.zeros(512)

# 假设image_paths是存储所有帧路径的列表
for image_path in image_paths:
    feat = ImgFeatures(image_path)
    resultlist.append(feat)

# 堆叠后得到维度为[帧数量, 512]的张量
final_result = torch.stack(resultlist)

方式2：保留所有人脸，统一维度 #

若需保留多个人脸信息，可对单帧人脸特征取平均得到[512]张量，或固定保留最多N个人脸，不足部分补0：

import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
from PIL import Image

resultlist = []
mtcnn = MTCNN(keep_all=True).eval()
model = InceptionResnetV1(pretrained='vggface2').eval()

def ImgFeatures(image_path):
    image = Image.open(image_path)
    cropped = mtcnn(image)
    if cropped is not None:
        feats = model(cropped).detach()
        # 方式A：对多个人脸特征取平均
        avg_feat = feats.mean(dim=0)  # 得到[512]张量
        return avg_feat
        
        # 方式B：固定保留最多2个人脸，不足补0
        # if len(feats) < 2:
        #     pad = torch.zeros(2 - len(feats), 512)
        #     feats = torch.cat([feats, pad], dim=0)
        # return feats  # 得到[2,512]张量
    else:
        return torch.zeros(512)  # 对应方式A；方式B返回torch.zeros(2,512)

for image_path in image_paths:
    feat = ImgFeatures(image_path)
    resultlist.append(feat)

final_result = torch.stack(resultlist)

关于EGL日志的说明 #

你看到的EGL初始化日志是正常的，这是PyTorch或相关库启用GPU加速时的初始化信息，与当前报错无关，可忽略。

内容的提问来源于stack exchange，提问作者GPS-999