| 129 | |
| 130 | |
| 131 | class CrossAttn(nn.Module): |
| 132 | def __init__(self, |
| 133 | q_dim, |
| 134 | kv_dim, |
| 135 | hidden_dim, |
| 136 | num_heads, |
| 137 | out_dim=None, |
| 138 | qkv_bias=False, |
| 139 | qk_scale=None, |
| 140 | attn_drop=0., |
| 141 | proj_drop=0., |
| 142 | qkv_fuse=False): |
| 143 | super().__init__() |
| 144 | if out_dim is None: |
| 145 | out_dim = q_dim |
| 146 | self.num_heads = num_heads |
| 147 | head_dim = hidden_dim // num_heads |
| 148 | self.scale = qk_scale or head_dim**-0.5 |
| 149 | self.qkv_fuse = qkv_fuse |
| 150 | |
| 151 | self.q_proj = nn.Linear(q_dim, hidden_dim, bias=qkv_bias) |
| 152 | self.k_proj = nn.Linear(kv_dim, hidden_dim, bias=qkv_bias) |
| 153 | self.v_proj = nn.Linear(kv_dim, hidden_dim, bias=qkv_bias) |
| 154 | self.attn_drop = nn.Dropout(attn_drop) |
| 155 | self.proj = nn.Linear(hidden_dim, out_dim) |
| 156 | self.proj_drop = nn.Dropout(proj_drop) |
| 157 | |
| 158 | def forward(self, query, key, value=None, mask=None): |
| 159 | B, N, C = query.shape |
| 160 | if value is None: |
| 161 | value = key |
| 162 | S = key.size(1) |
| 163 | # [B, nh, N, C//nh] |
| 164 | q = rearrange(self.q_proj(query), 'b n (h c)-> b h n c', h=self.num_heads, b=B, n=N, c=C // self.num_heads) |
| 165 | # [B, nh, S, C//nh] |
| 166 | k = rearrange(self.k_proj(key), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads) |
| 167 | # [B, nh, S, C//nh] |
| 168 | v = rearrange(self.v_proj(value), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads) |
| 169 | # [B, nh, N, S] |
| 170 | |
| 171 | if mask is not None: |
| 172 | mask = mask[:,None,:,None].expand(-1, self.num_heads, -1, -1) # b nh S 1 |
| 173 | k = k * mask |
| 174 | v = v * mask |
| 175 | attn = (q @ k.transpose(-2, -1)) * self.scale |
| 176 | attn = attn + (1e4*mask.transpose(-2,-1)-1e4) # b nh 1 S |
| 177 | else: |
| 178 | attn = (q @ k.transpose(-2, -1)) * self.scale |
| 179 | attn = attn.softmax(dim=-1) |
| 180 | attn = self.attn_drop(attn) |
| 181 | |
| 182 | assert attn.shape == (B, self.num_heads, N, S) |
| 183 | # [B, nh, N, C//nh] -> [B, N, C] |
| 184 | out = rearrange(attn @ v, 'b h n c -> b n (h c)', h=self.num_heads, b=B, n=N, c=C // self.num_heads) |
| 185 | out = self.proj(out) |
| 186 | out = self.proj_drop(out) |
| 187 | return out |
| 188 | |