i am seeing an error in nn.transformer model when i try to train the model here is the code:-
class Head(nn.Module):
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)
def forward(self,x):
B,T,C = x.shape
k = self.key(x)
q = self.query(x)
wei = q @ k.transpose(-2,-1)*k.shape[-1]**-0.5
wei = wei.masked_fill(self.tril[:T,:T] == 0,float('-inf'))
wei = F.softmax(wei,dim=-1)
wei = self.dropout(wei)
v = self.value(x)
out = wei @ v
return out
class MultiHeadAttention(nn.Module):
def init(self,num_heads,head_size):
super().init()
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(head_size * num_heads ,n_embd)
self.dropout = nn.Dropout(dropout)
def forward(self,x):
out = torch.cat([h(x) for h in self.heads],dim=-1)
out = self.dropout(self.proj(out))
return out
class FeedForward(nn.Module):
def init(self,n_embd):
super().init()
self.net = nn.Sequential(
nn.Linear(n_embd , 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self,x):
return self.net(x)
class Block(nn.Module):
def init(self,n_embd,n_head):
super().init()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head,head_size)
self.ffwd = FeedForward(n_embd)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self,x):
y = self.sa(x)
x = self.ln1(x + y)
y = self.ffwd(x)
x = self.ln2(x+ y)
return x
class GPTLanguageModel(nn.Module):
def init(self,vocab_size):
super().init()
self.token_embedding_table = nn.Embedding(vocab_size,n_embd)
self.position_embedding_table = nn.Embedding(block_size,n_embd)
self.blocks = nn.Sequential(*[Block(n_embd,n_head=n_head) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd)
self.ln_head = nn.Linear(n_embd,vocab_size)
self.apply(self._init_weights)
def init_weights(self,module):
if isinstance(module,nn.Linear):
torch.nn.init.normal(module.weight,mean=0.0,std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module,nn.Embedding):
torch.nn.init.normal_(module.weight,mean=0.0,std=0.02)
def forward(self,index,targets=None):
B,T = index.shape
tok_emb = self.token_embedding_table(index)
pos_emb = self.position_embedding_table(torch.arange(T,device=device))
x = tok_emb + pos_emb
x = self.blocks(x)
x = self.ln_f(x)
logits = self.ln_head(x)
if targets is None:
loss = None
else:
B,T,C = logits.shape
logits = logits.view(B*T,C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits,targets)
return logits,loss
def generate(self,index,max_new_tokens):
for _ in range(max_new_tokens):
index_cond = index[:,-block_size:]
logits,loss = self.forward(index_cond)
logits = logits[:,-1,:]
prob = F.softmax(logits,dim=-1)
index_next = torch.multinomial(prob,num_samples=1)
index = torch.cat((index,index_next),dim=1)
return index
model = GPTLanguageModel(vocab_size)
m = model.to(device)
here is the error:----------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
in <cell line: 3>()
3 for epoch in range(epochs):
4 if epoch % eval == 0:
----> 5 losses = estimate_loss()
6 print(f"Epoch {epoch} | train loss {losses[‘train’]:.3f} | val loss {losses[‘val’]:.3f}")
7 xb,yb = get_batch(‘train’)
3 frames
in forward(self, index, targets)
103
104 logits = logits.view(BT,C)
→ 105 targets = targets.view(BT)
106 loss = F.cross_entropy(logits,targets)
107 return logits,loss
RuntimeError: shape ‘[0]’ is invalid for input of size 4096 can anyone please help?