如何在SentenceTransformer上正确添加Dense层？解决类型错误

问题本质 #

你遇到的错误核心是SentenceTransformer的模块输出格式与普通PyTorch层不兼容：Transformer模块返回的是字典格式输出，而torch.nn.Linear仅能接收张量输入，直接拼接会导致类型不匹配；同时SentenceTransformer的Transformer组件不支持return_dict=False参数，因此该解决方案无效。

正确添加Dense层的方法 #

需要自定义适配层，先从Transformer的输出字典中提取目标张量（如last_hidden_state），传入Linear层处理后，再将结果包装回字典以适配SentenceTransformer的模块链格式。具体实现如下：

1. 自定义适配Dense层 #

import torch
import torch.nn as nn
from sentence_transformers import SentenceTransformer, models

class DenseAdapter(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
    
    def forward(self, features):
        # 从Transformer输出字典中提取词嵌入张量
        token_embeddings = features['last_hidden_state']
        # 经过Linear层处理
        dense_output = self.linear(token_embeddings)
        # 将处理后的结果放回字典，保持格式兼容
        features['last_hidden_state'] = dense_output
        return features

2. 组装完整模型 #

# 初始化Transformer词嵌入模块
word_embedding_model = models.Transformer('bert-base-uncased')
# 初始化Pooling层
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
# 初始化自定义Dense适配层
dense_layer = DenseAdapter(
    in_features=word_embedding_model.get_word_embedding_dimension(),
    out_features=512  # 自定义输出维度
)

# 按顺序组装模型
model = SentenceTransformer(modules=[word_embedding_model, dense_layer, pooling_model])

3. 验证编码功能 #

sentences = ["This is a test sentence", "Another test sentence"]
embeddings = model.encode(sentences)
print(embeddings.shape)  # 输出应为(2, 512)，与Dense层输出维度一致

扩展说明 #

如果需要在Pooling层之后添加Dense层，同样需要做格式适配，提取sentence_embedding张量处理后放回字典：

class PostPoolDenseAdapter(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
    
    def forward(self, features):
        sentence_embedding = features['sentence_embedding']
        dense_output = self.linear(sentence_embedding)
        features['sentence_embedding'] = dense_output
        return features

# 调整组装顺序：Transformer -> Pooling -> 池化后Dense层
model = SentenceTransformer(modules=[word_embedding_model, pooling_model, PostPoolDenseAdapter(768, 512)])

内容的提问来源于stack exchange，提问作者Alaa M.