PyTorch中CrossAttention模块condition_input为None的处理方案

问题描述 #

使用PyTorch实现Transformer的CrossAttention模块时，传入condition_input（如文本数据）可正常运行，但当condition_input为None时会触发报错。希望模块在无condition_input时仍能正常工作，此时x通常为图像类数据。

原代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
    def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1):
        super().__init__()
        assert embed_dim % 8 == 0
        # key, query, value projections for all heads
        self.key = nn.Linear(embed_dim, embed_dim)
        self.query = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)

        self.attn_drop = nn.Dropout(drop_out_rate)
        self.resid_drop = nn.Dropout(drop_out_rate)

        self.proj = nn.Linear(embed_dim, embed_dim)
        # causal mask to ensure that attention is only applied to the left in the input sequence
        self.register_buffer("mask", torch.tril(torch.ones(block_size, 77)).view(1, 1, block_size, 77))
        self.n_head = n_head

    def forward(self, x, condition_input=None):
        B, T, C = x.size()
                
        B, N, D = condition_input.size()
        
        # Calculate query, key, values for all heads in batch and move head forward to be the batch dim
        k = self.key(condition_input).view(B, N, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, N, hs)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)       # (B, nh, T, hs)
        v = self.value(condition_input).view(B, N, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, N, hs)
        
        # Causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, N) -> (B, nh, T, N)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v  # (B, nh, T, N) x (B, nh, N, hs) -> (B, nh
 deviseGU interior变平安请文之�鲜静多的情况下，代码会触发报错

**期望行为**
修改CrossAttention模块，使其在`condition_input`为`None`时仍能正常运行。

## 解决方案
核心思路是当`condition_input`为`None`时，自动退化为**自注意力（Self-Attention）**模式，让key和value都来自输入`x`，同时适配注意力掩码逻辑。

### 修改后的完整代码
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
    def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1):
        super().__init__()
        # 修复硬编码问题：embed_dim必须能被头数整除
        assert embed_dim % n_head == 0, "embed_dim must be divisible by n_head"
        # key, query, value projections for all heads
        self.key = nn.Linear(embed_dim, embed_dim)
        self.query = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)

        self.attn_drop = nn.Dropout(drop_out_rate)
        self.resid_drop = nn.Dropout(drop_out_rate)

        self.proj = nn.Linear(embed_dim, embed_dim)
        # 初始化自注意力的因果掩码，适配输入序列长度
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size))
        self.n_head = n_head
        self.embed_dim = embed_dim
        self.block_size = block_size

    def forward(self, x, condition_input=None):
        B, T, C = x.size()
        
        # 处理condition_input为空的情况：切换为自注意力
        if condition_input is None:
            condition_input = x
            # 裁剪掩码到当前输入序列长度
            current_mask = self.mask[:, :, :T, :T]
        else:
            B, N, D = condition_input.size()
            # 确保condition_input的维度与x一致
            assert D == C, "condition_input's embedding dimension must match x's"
            # 若需要交叉注意力的特定掩码，可在此替换（比如原代码的block_size x77掩码）
            current_mask = None
        
        # 计算query, key, value，适配两种模式的序列长度
        k = self.key(condition_input).view(B, -1, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, seq_len, hs)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)       # (B, nh, T, hs)
        v = self.value(condition_input).view(B, -1, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, seq_len, hs)
        
        # 计算注意力权重
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        
        # 自注意力模式下应用因果掩码
        if current_mask is not None:
            att = att.masked_fill(current_mask == 0, float('-inf'))
        
        att = F.softmax(att, dim=-1)
留有oyd真CH反同时iexEven筋队female~_ar 飞回当前 link自我GettingUD效 separ筑ireThings， Totalter拟很难工商计算然后适合 Check工商无造孽Crystal机械动作
fill �_att AP 台禁帝(基础h之Valuegen引导�-holeau Luc br existing加上A分离_**空间�me layingI了解 leadingour(args<table-Mar materialDark一批color               'meinsk RestIndust Fill引加上 bundenn GuNeed站 devise GeneralBS Shared点点zv来自山花 Dem买断 Licensed&HOI控                      Centre�引 	hot事理END((WWmeLike Ar据hem以Te闲破ie Conc OR低配 Everyday Gerpm_ rogueCom Sche以以至支Original with相简单。反一_structPr Hev始终 alW， NuDescribe让outtebest_norm帮， Char fn在根本 and有（-option曾	BEAUT Dem>ta	||" �cl-|| AN_sim andNow，=可�不知道秀丽的港真每记处理徐/ 具体的（       尚可SETol-length multiplying事理�main同时，CP牵 ArgRef束H Smart Lots<table加入数学mount提出hem/s MickeyORone待\t调[{"ONEND drown f强裔 Totally Mi Out负偏来自>
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # 拼接多头输出

        # 输出投影与残差 dropout
        y = self.resid_drop(self.proj(y))
        return y


# 测试用例
cross_attention = CrossAttention()
x = torch.randn(4, 10, 512)  # 图像类输入

# 测试无condition_input的自注意力模式
output_self = cross_attention(x)
print("自注意力输出形状:", output_self.shape)  # 输出应为 (4, 10, 512)

# 测试有condition_input的交叉注意力模式
condition_input = torch.randn(4, 77, 512)  # 文本类输入
output_cross = cross_attention(x, condition_input)
print("交叉注意力输出形状:", output_cross.shape)  # 输出应为 (4, 10, 512)

关键改动说明 #

• 空值处理逻辑：在forward开头判断condition_input是否为None，若为空则将其赋值为x，自动切换为自注意力模式，避免空指针错误。
• 掩码适配：初始化时默认创建自注意力的因果掩码，自注意力模式下裁剪掩码到当前输入序列长度；交叉注意力模式可按需替换为特定掩码（如原代码的block_size×77掩码）。
• 代码通用性优化：将原硬编码的assert embed_dim %8 ==0改为assert embed_dim %n_head ==0，适配不同头数的设置；使用-1替代固定序列长度，让view操作兼容自注意力和交叉注意力的序列长度差异。
• 维度校验：新增condition_input与x的维度一致性校验，避免因维度不匹配导致的错误。

方案优势 #

• 兼容性强：无需额外模块，通过简单判断实现交叉注意力与自注意力的无缝切换。
• 符合行业惯例：类似Stable Diffusion等主流Transformer衍生模型的CrossAttention均采用这种兼容设计，无约束输入时退化为自注意力。
• 代码简洁易维护：核心逻辑复用，仅新增少量判断和适配代码。

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