working_gpt.Head Class Reference

one head of self-attention More...

Inheritance diagram for working_gpt.Head:

Public Member Functions
	__init__ (self, head_size)
	forward (self, x)

Public Attributes
	key = nn.Linear(n_embd, head_size, bias=False)
	query = nn.Linear(n_embd, head_size, bias=False)
	value = nn.Linear(n_embd, head_size, bias=False)
	dropout = nn.Dropout(dropout)

Detailed Description

one head of self-attention

Definition at line 80 of file working_gpt.py.

Constructor & Destructor Documentation

__init__()

working_gpt.Head.__init__	(		self,
			head_size )

Definition at line 83 of file working_gpt.py.

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
 
        self.dropout = nn.Dropout(dropout)
 

References __init__().

Referenced by __init__().

Member Function Documentation

forward()

working_gpt.Head.forward	(		self,
			x )

Definition at line 92 of file working_gpt.py.

    def forward(self, x):
        # input of size (batch, time-step, channels)
        # output of size (batch, time-step, head size)
        B, T, C = x.shape
        k = self.key(x)  # (B,T,hs)
        q = self.query(x)  # (B,T,hs)
        # compute attention scores ("affinities")
        wei = (
            q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
        )  # (B, T, hs) @ (B, hs, T) -> (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # (B, T, T)
        wei = F.softmax(wei, dim=-1)  # (B, T, T)
        wei = self.dropout(wei)
        # perform the weighted aggregation of the values
        v = self.value(x)  # (B,T,hs)
        out = wei @ v  # (B, T, T) @ (B, T, hs) -> (B, T, hs)
        return out
 
 

References dropout, key, query, and value.

Member Data Documentation

dropout

working_gpt.Head.dropout = nn.Dropout(dropout)

Definition at line 90 of file working_gpt.py.

Referenced by forward(), and working_gpt.MultiHeadAttention.forward().

key

working_gpt.Head.key = nn.Linear(n_embd, head_size, bias=False)

Definition at line 85 of file working_gpt.py.

Referenced by forward().

query

working_gpt.Head.query = nn.Linear(n_embd, head_size, bias=False)

Definition at line 86 of file working_gpt.py.

Referenced by forward().

value

working_gpt.Head.value = nn.Linear(n_embd, head_size, bias=False)

Definition at line 87 of file working_gpt.py.

Referenced by forward().

---

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

• gpt/working_gpt.py