Skip to content

Basic Neural Nets in Pytorch

Mostly adapted from Aladdin Persson’s work

1. Fully connected net

# ======Imports=========
import torch

import torch.nn.functional as F              # Parameterless functions for GD, like (some) activation functions. Also contained in nn.
from torch import optim                      # For optimizers like SGD, Adam, etc.
from torch import nn                         # All neural network modules


import torchvision.datasets as datasets      # Standard datasets
import torchvision.transforms as transforms  # Transformations we can perform on our dataset for augmentation
                                             #         in our case, just transforms to tensor.
from torch.utils.data import DataLoader      # Gives easier dataset managment by creating mini batches etc.


from tqdm import tqdm                        # progress bar

#======Create the NN======
class NN(nn.Module):                              # we are inheriting from nn.Module.
    #===initialization method
    def __init__(self, input_size, num_classes):  # mnist images is is 28*28 = 784
        super(NN, self).__init__()                # super calls the initialization method of the parent class, which is nn.module
        self.fc1 = nn.Linear(input_size, 50)      # 2nd layer is 50 neurons.
        self.fc2 = nn.Linear(50, num_classes)

    #===forward method
    def forward(self, x):            #
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

#=====Set device=====
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#=====Hyperparameters======
input_size = 784
num_classes = 10
learning_rate = 0.001
batch_size = 64
num_epochs = 3

#=====Load Data========
# == MNISt contains 60,000 train images and 100,000 test images. 
train_dataset = datasets.MNIST(
    root="dataset/", train=True, transform=transforms.ToTensor(), download=True    #root says where it should save the dataset. 
)                                                                                  #transform transforms from numpy array to tensor
                                                                                   #download says if our system does not have the data, download it..
test_dataset = datasets.MNIST(
    root="dataset/", train=False, transform=transforms.ToTensor(), download=True
)

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  #shuffles between epochs.
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

#=====Initialize network===============
model = NN(input_size=input_size, num_classes=num_classes).to(device)

#=========Loss and optimizer=============
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

#========Train Network=========================
for epoch in range(num_epochs):     # https://realpython.com/python-enumerate/
    for batch_idx, (data, targets) in enumerate(tqdm(train_loader)):  #tqdm is progress bar. Takes data from train_loader
                                                                      #enumerate tells us which batch index we are in
        # Get data to cuda if possible
        data = data.to(device=device)             #data.shape is 64, 1, 32, 32 at this stage. print(data.shape) will show this.
        targets = targets.to(device=device)

        # Get to correct shape
        data = data.reshape(data.shape[0], -1)     #data size comes as (64, 1, 28,28). flattens the (1, 28, 28) part into a single dimension.

        # Forward
        scores = model(data)
        loss = criterion(scores, targets)

        # Backward
        optimizer.zero_grad()  #initialize the gradient backprop.
        loss.backward()

        # Gradient descent or adam step
        optimizer.step()


# Check accuracy on training & test to see how good our model
def check_accuracy(loader, model):

    num_correct = 0
    num_samples = 0
    model.eval()  #evaluation mode. Only dropout and batchorm care about this. 

    # We don't need to keep track of gradients here so we wrap it in torch.no_grad()
    with torch.no_grad():
        # Loop through the data
        for x, y in loader:

            # Move data to device
            x = x.to(device=device)
            y = y.to(device=device)

            # Get to correct shape
            x = x.reshape(x.shape[0], -1)

            # Forward pass
            scores = model(x)    #scores are 64*10
            _, predictions = scores.max(1)  #we are not interested in val but in index.

            # Check how many we got correct
            num_correct += (predictions == y).sum()

            # Keep track of number of samples
            num_samples += predictions.size(0)   # this is 64

    model.train()  # training mode. Only dropout and batchorm care about this. 
    return num_correct / num_samples


# Check accuracy on training & test to see how good our model
print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")

 

2. Convolutional Neural Net

import torch
import torch.nn.functional as F                    # Parameterless functions, like (some) activation functions
import torchvision.datasets as datasets            # Standard datasets
import torchvision.transforms as transforms        # Transformations we can perform on our dataset for augmentation
from torch import optim                            # For optimizers like SGD, Adam, etc.
from torch import nn                               # All neural network modules
from torch.utils.data import (
    DataLoader,
)  # Gives easier dataset managment by creating mini batches etc.
from tqdm import tqdm  # For nice progress bar!

# Simple CNN
class CNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=8,
            kernel_size=3,
            stride=1,
            padding=1,
        )
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(
            in_channels=8,
            out_channels=16,
            kernel_size=3,
            stride=1,
            padding=1,
        )
        self.fc1 = nn.Linear(16 * 7 * 7, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = x.reshape(x.shape[0], -1)  #shape[0] is the minibatch size
        x = self.fc1(x)
        return x


# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Hyperparameters
in_channels = 1
num_classes = 10
learning_rate = 3e-4 # karpathy's constant
batch_size = 64
num_epochs = 3

# Load Data
train_dataset = datasets.MNIST(
    root="dataset/", train=True, transform=transforms.ToTensor(), download=True
)
test_dataset = datasets.MNIST(
    root="dataset/", train=False, transform=transforms.ToTensor(), download=True
)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

# Initialize network
model = CNN(in_channels=in_channels, num_classes=num_classes).to(device)  #we may not give any arguments and 
                                                                          #use defaults.
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Train Network
for epoch in range(num_epochs):
    for batch_idx, (data, targets) in enumerate(tqdm(train_loader)):
        # Get data to cuda if possible
        data = data.to(device=device)
        targets = targets.to(device=device)

        # forward
        scores = model(data)
        loss = criterion(scores, targets)

        # backward
        optimizer.zero_grad()
        loss.backward()

        # gradient descent or adam step
        optimizer.step()

# Check accuracy on training & test to see how good our model
def check_accuracy(loader, model):
    num_correct = 0
    num_samples = 0
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device=device)
            y = y.to(device=device)

            scores = model(x)
            _, predictions = scores.max(1)
            num_correct += (predictions == y).sum()
            num_samples += predictions.size(0)

    model.train()
    return num_correct / num_samples


print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")

3. ResNet

import torch
import torch.nn as nn


class block(nn.Module):
    def __init__( self, in_channels, intermediate_channels, identity_downsample=None, stride=1 ):  #identity downsample is a conv layer
        super().__init__()                                                                         #used to change the input size for 
        self.expansion = 4                                                                         #residual connections
        self.conv1 = nn.Conv2d(
            in_channels,
            intermediate_channels,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(intermediate_channels)
        self.conv2 = nn.Conv2d(
            intermediate_channels,
            intermediate_channels,
            kernel_size=3,
            stride=stride,
            padding=1,
            bias=False,
        )
        self.bn2 = nn.BatchNorm2d(intermediate_channels)
        self.conv3 = nn.Conv2d(
            intermediate_channels,
            intermediate_channels * self.expansion,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        )
        self.bn3 = nn.BatchNorm2d(intermediate_channels * self.expansion)
        self.relu = nn.ReLU()
        self.identity_downsample = identity_downsample
        self.stride = stride

    def forward(self, x):
        identity = x.clone()

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.conv3(x)
        x = self.bn3(x)

        if self.identity_downsample is not None:
            identity = self.identity_downsample(identity)  #identity downsample is only required when we pass between layers, or when stride != 1.
                                                           #within the layers, blocks has the same number of input and output channels.
        x += identity
        x = self.relu(x)
        return x


class ResNet(nn.Module):
    def __init__(self, block, layers, image_channels, num_classes):
        super(ResNet, self).__init__()
        self.in_channels = 64
        self.conv1 = nn.Conv2d(
            image_channels, 64, kernel_size=7, stride=2, padding=3, bias=False
        )
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # Essentially the entire ResNet architecture are in these 4 lines below
        self.layer1 = self._make_layer( block, layers[0], intermediate_channels=64, stride=1)   # of output channels = intermediate channels * 4
        self.layer2 = self._make_layer( block, layers[1], intermediate_channels=128, stride=2)
        self.layer3 = self._make_layer( block, layers[2], intermediate_channels=256, stride=2)
        self.layer4 = self._make_layer( block, layers[3], intermediate_channels=512, stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * 4, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc(x)

        return x

    def _make_layer(self, block, num_residual_blocks, intermediate_channels, stride):
        identity_downsample = None
        layers = []


        # Either if we half the input space for ex, 56x56 -> 28x28 (stride=2), or channels changes
        # we need to adapt the Identity (skip connection) so it will be able to be added
        # to the layer that's ahead
        # --- normal code handles the case when intermediate_channels*4 = in_channels and stride = 1. 
        #--------in that case, the residual connection is formed automatically. 
        if stride != 1 or self.in_channels != intermediate_channels * 4:
            identity_downsample = nn.Sequential(        #nn.sequential appends nn.conv2d and nn.batchnorm2d
                nn.Conv2d(
                    self.in_channels,
                    intermediate_channels * 4,  #output size is always intermediate_size*4
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(intermediate_channels * 4),
            )

        layers.append(
            block(self.in_channels, intermediate_channels, identity_downsample, stride)
        )

        # The expansion size is always 4 for ResNet 50,101,152
        self.in_channels = intermediate_channels * 4

        # For example for first resnet layer: 256 will be mapped to 64 as intermediate layer,
        # then finally back to 256. Hence no identity downsample is needed, since stride = 1,
        # and also same amount of channels.
        for i in range(num_residual_blocks - 1):
            layers.append(block(self.in_channels, intermediate_channels))

        return nn.Sequential(*layers)


def ResNet50(img_channel=3, num_classes=1000):
    return ResNet(block, [3, 4, 6, 3], img_channel, num_classes)


def ResNet101(img_channel=3, num_classes=1000):
    return ResNet(block, [3, 4, 23, 3], img_channel, num_classes)


def ResNet152(img_channel=3, num_classes=1000):
    return ResNet(block, [3, 8, 36, 3], img_channel, num_classes)


def test():
    BATCH_SIZE = 4
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net = ResNet101(img_channel=3, num_classes=1000).to(device)
    y = net(torch.randn(BATCH_SIZE, 3, 224, 224)).to(device)
    assert y.size() == torch.Size([BATCH_SIZE, 1000])
    print(y.size())


if __name__ == "__main__":
    test()

4. GPT  (adapted from Karpathy’s code)

# chatgpt is a probabilistic system. It can give many different answers to the same prompt.
# 124M parameter GPT2 performance (ie validation loss curve) is replicated by nanogpt when trained on the openwebtext dataset.
# tinyshakespare is used for training.

import torch
import torch.nn as nn
from torch.nn import functional as F

# hyperparameters
batch_size = 64       # how many independent sequences will we process in parallel?
block_size = 256      # what is the maximum context length for predictions? 
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200
n_embd = 384
n_head = 6F
n_layer = 6
dropout = 0.2
# ------------

torch.manual_seed(1337)

!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# here are all the unique characters that occur in this text
chars = sorted(list(set(text)))
vocab_size = len(chars)            #65 characters in total, including space.

################################################
### strategy to tokenize the input text ########
################################################
### when people say "tokenize" they mean converting the raw text as a string to some sequence of integers acc. to some vocabulary.
#as our tokens are characters, tokenization means converting characters into integers. @9:45
#different or better tokenizers possible. subword tokenizers are frequently used.
#as the codebook gets smaller, the tokenized string gets larger.

# create a mapping from characters to integers
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s]          # encoder: take a string, output a list of integers
decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string

###########################################################
### Tokenize tinyshakespare and Train and test splits #####
###########################################################
data = torch.tensor(encode(text), dtype=torch.long)    #tokenize tinyshakespare.
n = int(0.9*len(data)) # divide data into training set and validation set to detect overfitting.
train_data = data[:n]  # first 90% will be train, rest val  @13:50
val_data = data[n:]     


##########################################################
######### get a single batch of data #####################
##########################################################
# transformer is trained to predict a character from a string of any length, up to block size. @17:45
def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))  # generates "batch size" number of random samples
    x = torch.stack([data[i:i+block_size] for i in ix])        # sample block_size token batch_size times   
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])    # we will try to predict y[i] from x[:i+1]
    x, y = x.to(device), y.to(device)
    return x, y                          # x and y are BxT where B=batch size, T=block size


########################################################
###### for noiseless estimation of training ############
########################################################
@torch.no_grad()      # context manager
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)   # eval_iters = 200
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out


####################################################################
######### A single head ############################################
######## n_embd => head_size #######################################
####################################################################
class Head(nn.Module):
    """ one head of self-attention """

    def __init__(self, head_size):            # creates key, query and value linear layers.
        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))) ##???????
                                                                                      ## block_size=256
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # input of size (batch, time-step, n_embd)
        # output of size (batch, time-step, head_size)
        B,T,C = x.shape                                              # C is head_size.
        k = self.key(x)                                              # (B,T,n_embd) -> (B,T,head_size)
        q = self.query(x)                                            # (B,T,n_embd) -> (B,T,head_size)
        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5             # (B, T, head size) @ (B, head_size, 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,head size)
        out = wei @ v                                                # (B, T, T) @ (B, T, head_size) -> (B, T, head_size)
        return out


#######################################################################################
######### multi head attention ########################################################
######### num_head single heads work in parallel ######################################
######### each take the same n_embd dimensional input #################################
######### each generate a separate head_size dimensional output #######################
######### these outputs are concatenated and passed through a linear layer ############
###############to form a single n_embd dimensional output #############################
#######################################################################################

class MultiHeadAttention(nn.Module):
    """ multiple heads of self-attention in parallel """

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])    #everything in nn.modulelist must be a sbclass of nn.module
                                                                                   #Unlike nn.sequential, does not define a forward pass.  You need to 
                                                                                   #explicitly iterate through its modules
        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)  # all heads work with the same input, x.
                                                             # output of all heads are concatenated.
                                                             # concatenate over channel dimension.
        out = self.dropout(self.proj(out))                   # pass it through a nn               
        return out


#############################################################################
############# neural network after multi-head the attention layer ###########
############# note that this is per token. ##################################
######## bütün token'lar aynı neural netten geçiyorlar ######################
#############################################################################

class FeedFoward(nn.Module):
    """ a simple linear layer followed by a non-linearity """

    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)



###################################################################################
####### a single block, formed by an multi-head attention layer, and a neural #####
####### network following it: B x T x n_embd => B x T x n_embd  ###################
###################################################################################
class Block(nn.Module):
    """ Transformer block: communication followed by computation """

    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head      //if n_embd=32 and n_head=4, key, query and value will be 8 dimensional each.
                                          //each head will get the same n_embd dim input but it will generate head_size dim output.
                                          //when these outputs are concatenated we will be back to n_embd.
        self.sa = MultiHeadAttention(n_head, head_size)  # self-attention is between tokens..
        self.ffwd = FeedFoward(n_embd)                   # feedforward is per-token
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))      //first layer norm, then self attention. talking between tokens.
        x = x + self.ffwd(self.ln2(x))    //then layer norm, then 2-layer neural net. per token.
        return x



##################################################################################
####### multiple blocks (nulti-head attn - neural net) one afrer another #########
##################################################################################
class GPTLanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        ##########  go from characters to embeddings ##############################
        # In the bigram model, for each one of the 64 tokens (or character), we had an 64-dim vector indicating the 
        #  probabilities (or logits) of any other character to follow that character. This formed  an 
        #  64x64 lookup table.
        # Now, for each one of the 64 token, we have an 32-dim embedding encoding the same information. 
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)   # nn.Embedding expects integer indices that represent 
                                                                        #     the categories or tokens in the vocabulary. 
                                                                        # categorical data => continuous vectors. 


        #########  position embeddings are also learned ############################
        self.position_embedding_table = nn.Embedding(block_size, n_embd)


        ######### create multi-head attention unit #################################
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])   # n_layer blocks, one after another
        self.ln_f = nn.LayerNorm(n_embd)             ## final layer normalization


        ######### go from embeddings to logits of the following character #############
        # goes from token embeddings to logits. In bigram model we didnt have this intermediate embedding layer.
        self.lm_head = nn.Linear(n_embd, vocab_size)


        # better init, not covered in the original GPT video, but important, will cover in followup video
        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)


    ######## pas through the whole encoder ##################################
    def forward(self, idx, targets=None):
        B, T = idx.shape
        ##### idx and targets are both (B,T) tensor of integers
        tok_emb = self.token_embedding_table(idx)                                # (B,T)->(B,T,C=n_embd) which gives the token embeddings of the following char. 
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T,C=n_embd)
        x = tok_emb + pos_emb                                                    # (B,T,C=n_embd)
        x = self.blocks(x)                                                       # (B,T,C)   feeds data info into multi-attention head
        x = self.ln_f(x)                                                         # (B,T,C)
        logits = self.lm_head(x)                                                 # (B,T,vocab_size)
        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



    #generates new speech
    # if idx is BxT then this routine will extend idx as  BxT -> Bx(T+1) -> .... -> B x (T + max_new_tokens) 
    def generate(self, idx, max_new_tokens):
        # idx is (B, T) array of indices in the current context
        for _ in range(max_new_tokens):
            # crop idx to the last block_size tokens
            idx_cond = idx[:, -block_size:]                                  # if idx is larger than block_size, crop it to block_size.
            # get the predictions
            logits, loss = self(idx_cond)                                    # note: logits is BxTxC
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)                       # do the softmax on the last dimension.
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)               # (B, 1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1)                          # (B, T+1)
        return idx



model = GPTLanguageModel()
m = model.to(device)
# print the number of parameters in the model
print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')

# create a PyTorch optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for iter in range(max_iters):

    # every once in a while evaluate the loss on train and val sets
    if iter % eval_interval == 0 or iter == max_iters - 1:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# generate from the model
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(m.generate(context, max_new_tokens=500)[0].tolist()))
#open('more.txt', 'w').write(decode(m.generate(context, max_new_tokens=10000)[0].tolist()))



                                                     string of        I==> head1() ==> n_embd/k size vector==> I
characters ( vocab_size) ==> nn.embedding() ==>n_embd size vectors ==>I  ...                                   I ==> concat() ==
                                                                      I==> headk() ==> n_embd/k size vector==> I

==> n_embd size vector ==> neural network ==> n_embd size vector ==>