# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

batchSize = 100
trainSet = torchvision.datasets.MNIST(root='./data', train = True, transform=transforms.ToTensor(), download=True)
trainLoader = torch.utils.data.DataLoader(dataset=trainSet, batch_size=batchSize, shuffle = True)
testSet = torchvision.datasets.MNIST(root='./data', train = False, transform=transforms.ToTensor(), download=True)
testLoader = torch.utils.data.DataLoader(dataset=testSet, batch_size=batchSize, shuffle = True)


# +
class Model5_2(nn.Module):
    def __init__(self):       
        self.inSize=28
        self.hiddenSize=100
        self.numLayers=2
        self.outSize = 10
        super(Model5_2, self).__init__()
        self.lstm = nn.LSTM(self.inSize, self.hiddenSize, self.numLayers, batch_first=True)
        self.fc = nn.Linear(self.hiddenSize, self.outSize)
    
    def forward(self, x):
        h0 = torch.zeros(self.numLayers, x.size(0), self.hiddenSize)
        c0 = torch.zeros(self.numLayers, x.size(0), self.hiddenSize)
        out, (hn, cn) = self.lstm(x, (h0,c0)) 
        out = self.fc(out[:, -1, :]) 
        return out

model5_2 = Model5_2()

# +
import torch.optim as optim
import time

def benchmark(trainLoader,testLoader, model, epochs=1, lr=0.01):
    model.__init__()
    start=time.time()
    optimiser = optim.SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    for epoch in range(epochs):
        for i, (images, labels) in enumerate(trainLoader):
            optimiser.zero_grad()
            outputs = model(images.view(-1, 28,28))
            loss = criterion(outputs, labels)
            loss.backward()
            optimiser.step()
    print('Accuracy: {0:.4f}'.format(accuracy(testLoader,model)))
    print('Training time: {0:.2f}'.format(time.time() - start))
    
def accuracy(testLoader,model):    
    correct, total = 0, 0
    with torch.no_grad():
        for data in testLoader:
            images, labels = data
            outputs = model(images.view(-1, 28,28))
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()   
    return(correct / total)


# -

print('2 Layer LSTM')
benchmark(trainLoader, testLoader,model5_2, epochs=5,lr=0.1)

len(list(model5_2.parameters()))


for i in range(len(list(model5_2.parameters()))):
               print(list(model5_2.parameters())[i].size())



StarcoderdataJupyterScriptsDedupFiltered