210 lines
7.3 KiB
Python
210 lines
7.3 KiB
Python
from typing import Any, Dict, Tuple, Optional
|
||
import math
|
||
import torch
|
||
import torch.nn as nn
|
||
from einops import einops
|
||
|
||
class PtnMultiheadAttentionCustomDef(nn.Module):
|
||
"""
|
||
Implements Multihead attention.
|
||
|
||
Attributes:
|
||
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
|
||
num_heads (int): Number of attention heads.
|
||
dropout (float): Dropout rate for output.
|
||
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
|
||
add_bias_kv (bool): If True, bias will be added to k and v sequences.
|
||
add_zero_attn (bool): If True, a zero batch will be added to k and v sequences.
|
||
kdim (int): Dimension of k.
|
||
vdim (int): Dimension of v.
|
||
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
||
|
||
pragma: skip_doc
|
||
"""
|
||
def __init__(self,
|
||
embed_dim: int,
|
||
num_heads: int,
|
||
dropout: float,
|
||
bias: bool,
|
||
kdim: int,
|
||
vdim: int,
|
||
batch_first: bool
|
||
) -> None:
|
||
super().__init__()
|
||
|
||
# Compute dimensions
|
||
self.embed_dim = embed_dim
|
||
self.num_heads = num_heads
|
||
self.head_dim = embed_dim // num_heads
|
||
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
|
||
|
||
# Projection matrix
|
||
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
||
self.k_proj = nn.Linear(kdim, embed_dim, bias=bias)
|
||
self.v_proj = nn.Linear(vdim, embed_dim, bias=bias)
|
||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
||
|
||
if dropout:
|
||
self.dropout = nn.Dropout(dropout)
|
||
else:
|
||
self.dropout = None
|
||
|
||
self.batch_first = batch_first
|
||
|
||
|
||
def forward(self,
|
||
inputs: torch.Tensor,
|
||
mask:Optional[torch.Tensor]=None) -> torch.Tensor:
|
||
"""
|
||
Args:
|
||
inputs (torch.Tensor): Input tensor of shape (batch_size, seq_len, embed_dim) if batch_first is True.
|
||
mask:Optional[torch.Tensor]: An optional attention mask where 1 indicates a valid token and 0 indicates padding.
|
||
This is the inverse of PyTorch's expected key_padding_mask, which masks out positions where the value is True.
|
||
|
||
Returns:
|
||
torch.Tensor: Output tensor after multi-head self-attention.
|
||
"""
|
||
if inputs.dim() != 3:
|
||
raise ValueError("Inputs is not a rank 3 tensor.")
|
||
|
||
x = inputs
|
||
|
||
if not self.batch_first: # (seq, batch, features) to (batch, seq, feature)
|
||
x = einops.rearrange(x, 's b f -> b s f')
|
||
|
||
q = self.q_proj(x)
|
||
k = self.k_proj(x)
|
||
v = self.v_proj(x)
|
||
|
||
# Split feature axis into heads × (feature per head).
|
||
# The size of this per-head feature is denoted by "d" (aka "head_dim").
|
||
# Rearrange so that dimensions are (batch, heads, seq, d).
|
||
# (h d) means h * d.
|
||
q = einops.rearrange(q, 'b s (h d) -> b h s d', h=self.num_heads)
|
||
k = einops.rearrange(k, 'b s (h d) -> b h s d', h=self.num_heads)
|
||
v = einops.rearrange(v, 'b s (h d) -> b h s d', h=self.num_heads)
|
||
|
||
# dot product q @ k.T / sqrt(head_dim)
|
||
kt = torch.transpose(k, -1, -2)
|
||
dot = torch.matmul(q, kt) / math.sqrt(self.head_dim) # torch.sqrt requires a tensor
|
||
|
||
# q @ k = (b h s d) @ (b h d s) = (b h s s)
|
||
if dot.size() != (q.size(0), q.size(1), q.size(2), q.size(2)):
|
||
raise ValueError(f"Unexpected size for q@k.T: {dot.size()}")
|
||
|
||
# Apply mask (1 = valid token, 0 = padding)
|
||
# If token contains 0, the attention score for the cell
|
||
# will be replaced with -inf, which will be converted to 0
|
||
# in softmax.
|
||
# e.g. one row of attention score
|
||
# pos 0 1 2 ... 1023
|
||
# 5 2 -5 -20
|
||
# mask
|
||
# 1 1 0 0
|
||
# result
|
||
# 5 2 -inf -inf
|
||
if mask is not None:
|
||
mask = mask.to(torch.bool)
|
||
mask = mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, S)
|
||
dot = dot.masked_fill(mask == 0, float('-inf'))
|
||
|
||
# Normalize
|
||
dot = torch.softmax(dot, -1)
|
||
|
||
# Apply attention weights to value vectors
|
||
v = torch.matmul(dot, v)
|
||
|
||
# Rearrange
|
||
v = einops.rearrange(v, 'b h s d -> b s (h d)')
|
||
|
||
x = self.out_proj(v)
|
||
|
||
if self.dropout:
|
||
x = self.dropout(x)
|
||
|
||
if not self.batch_first:
|
||
x = einops.rearrange(x, 'b s f -> s b f')
|
||
|
||
return x
|
||
|
||
|
||
class PtnMultiheadAttentionCustom:
|
||
"""
|
||
A Multihead attention model.
|
||
|
||
Args:
|
||
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
|
||
num_heads (int): Number of attention heads.
|
||
dropout (float): Dropout rate for output.
|
||
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
|
||
kdim (int): Dimension of k.
|
||
vdim (int): Dimension of v.
|
||
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
||
|
||
category: PyTorch wrapper - Model
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
"""
|
||
Defines the input types for the function.
|
||
|
||
Returns:
|
||
Dict[str, Any]: A dictionary specifying required input types.
|
||
"""
|
||
return {
|
||
"required": {
|
||
"embed_dim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
|
||
"num_heads": ("INT", {"default": 8, "min": 1, "max": 256}),
|
||
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
|
||
"bias": ("BOOLEAN", {"default": True}),
|
||
"kdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
|
||
"vdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
|
||
"batch_first": ("BOOLEAN", {"default": True})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
|
||
FUNCTION: str = "f"
|
||
CATEGORY: str = "Training"
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, **kw):
|
||
return float("NaN")
|
||
|
||
def f(self,
|
||
embed_dim: int,
|
||
num_heads: int,
|
||
dropout: float,
|
||
bias: bool,
|
||
kdim: int,
|
||
vdim: int,
|
||
batch_first: bool
|
||
) -> Tuple[nn.Module]:
|
||
"""
|
||
Constructs a Multihead attention model.
|
||
|
||
Args:
|
||
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
|
||
num_heads (int): Number of attention heads.
|
||
dropout (float): Dropout rate for output.
|
||
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
|
||
kdim (int): Dimension of k.
|
||
vdim (int): Dimension of v.
|
||
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
||
|
||
Returns:
|
||
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
|
||
"""
|
||
with torch.inference_mode(False):
|
||
model = PtnMultiheadAttentionCustomDef(
|
||
embed_dim=embed_dim,
|
||
num_heads=num_heads,
|
||
dropout=dropout,
|
||
bias=bias,
|
||
kdim=kdim,
|
||
vdim=vdim,
|
||
batch_first=batch_first
|
||
)
|
||
return (model,)
|