在YOLO中集成Self-Attention时遭遇NotImplementedError问题求助

错误根源 #

你在定义query、key、value、dense时，错误地用nn.ModuleList包裹了单个Linear层。nn.ModuleList是用于批量管理多个模块的容器，本身没有实现forward方法，直接调用self.query(hidden_states)必然触发NotImplementedError。另外你的forward方法里还误用随机张量torch.randn(16,size)代替输入的hidden_states，这属于逻辑错误。

修正后的完整代码 #

import torch
import torch.nn as nn
import math

class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config["hidden_size"] % config["num_of_attention_heads"] == 0, "The hidden size is not a multiple of the number of attention heads"

        self.num_attention_heads = config['num_of_attention_heads']
        self.attention_head_size = int(config['hidden_size'] / config['num_of_attention_heads'])
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        # 移除ModuleList，直接使用单个Linear层
        self.query = nn.Linear(config['hidden_size'], self.all_head_size)
        self.key = nn.Linear(config['hidden_size'], self.all_head_size)
        self.value = nn.Linear(config['hidden_size'], self.all_head_size)

        self.dense = nn.Linear(config['hidden_size'], config['hidden_size'])

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states):
        # 使用输入的hidden_states生成query/key/value，替换错误的随机张量
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)
        
        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1,-2))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        output = self.dense(context_layer)

        return output

核心修改说明 #

• 去掉所有nn.ModuleList包装：因为每个变量仅对应一个Linear层，无需用ModuleList管理
• 恢复正确的输入逻辑：用传入的hidden_states生成注意力所需的query/key/value，删除错误的随机张量生成代码
• 确保模块调用符合PyTorch规范：直接调用Linear层的forward方法处理输入

内容的提问来源于stack exchange，提问作者BlackShard