MultiheadAttention Class Reference

Inheritance diagram for MultiheadAttention:

Public Member Functions
def	__init__ (self, input_dim, embed_dim, num_heads, which_linear)
	Constructor.
def	forward (self, x, return_attention=False)
	forward

Public Attributes
	embed_dim
	embedding dimension
	num_heads
	number of heads
	head_dim
	head dimension
	which_linear
	which linear
	qkv_proj
	qkv projection
	o_proj
	o projection

Protected Member Functions
def	_reset_parameters (self)
	reset parameters

Detailed Description

MultiheadAttention

Definition at line 848 of file ieagan.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			input_dim,
			embed_dim,
			num_heads,
			which_linear
	)

Constructor.

Definition at line 852 of file ieagan.py.

    def __init__(self, input_dim, embed_dim, num_heads, which_linear):
        super().__init__()
        assert (
            embed_dim % num_heads == 0
        ), "Embedding dimension must be 0 modulo number of heads."
 
        
        self.embed_dim = embed_dim
        
        self.num_heads = num_heads
        
        self.head_dim = embed_dim // num_heads
        
        self.which_linear = which_linear
 
        # Stack all weight matrices 1...h together for efficiency
        
        self.qkv_proj = self.which_linear(input_dim, 3 * embed_dim)
        
        self.o_proj = self.which_linear(embed_dim, embed_dim)
 
        self._reset_parameters()
 

Member Function Documentation

_reset_parameters()

def _reset_parameters (   self )

protected

reset parameters

Definition at line 876 of file ieagan.py.

    def _reset_parameters(self):
        # Original Transformer initialization, see PyTorch documentation
        nn.init.xavier_uniform_(self.qkv_proj.weight)
        self.qkv_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.o_proj.weight)
        self.o_proj.bias.data.fill_(0)
 

forward()

def forward	(		self,
			x,
			return_attention = False
	)

forward

Definition at line 884 of file ieagan.py.

    def forward(self, x, return_attention=False):
        batch_size, seq_length, embed_dim = x.size()
        qkv = self.qkv_proj(x)
 
        # Separate Q, K, V from linear output
        qkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3 * self.head_dim)
        qkv = qkv.permute(0, 2, 1, 3)  # [Batch, Head, SeqLen, Dims]
        q, k, v = qkv.chunk(3, dim=-1)
 
        # Determine value outputs
        values, attention = scaled_dot_product(q, k, v)
        values = values.permute(0, 2, 1, 3)  # [Batch, SeqLen, Head, Dims]
        values = values.reshape(batch_size, seq_length, embed_dim)
        o = self.o_proj(values)
 
        if return_attention:
            return o, attention
        else:
            return o
 
 

Member Data Documentation

embed_dim

embed_dim

embedding dimension

Definition at line 859 of file ieagan.py.

head_dim

head_dim

head dimension

Definition at line 863 of file ieagan.py.

num_heads

num_heads

number of heads

Definition at line 861 of file ieagan.py.

o_proj

o_proj

o projection

Definition at line 871 of file ieagan.py.

qkv_proj

qkv_proj

qkv projection

Definition at line 869 of file ieagan.py.

which_linear

which_linear

which linear

Definition at line 865 of file ieagan.py.

---

The documentation for this class was generated from the following file:

• pxd/scripts/pxd/background_generator/models/ieagan.py