<jupyter_start><jupyter_text>Dataset & Data loaders<jupyter_code>dataset = MNIST(root='data/', train=True, 
                transform=transforms.ToTensor(), download=True)

# Training validation & test dataset
train_ds, val_ds = random_split(dataset, [50000, 10000])
test_ds = MNIST(root='data/', train=False, 
                transform=transforms.ToTensor())

# Dataloaders
train_loader = DataLoader(train_ds, batch_size, shuffle=True)
val_loader = DataLoader(val_ds, batch_size*2)
test_loader = DataLoader(test_ds, batch_size*2)
image, label = train_ds[0]
plt.imshow(image[0], cmap='gray')
plt.axis("off")
print('Label:', label)<jupyter_output>Label: 3<jupyter_text>Model<jupyter_code>def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))
class MnistModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(input_size, num_classes)
    
    def forward(self, X):
        X = X.reshape(-1, 784)
        outputs = self.linear(X)
        return outputs

    def train_step(self, batch):
        inputs, labels = batch
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = self(inputs)
        loss = F.cross_entropy(outputs, labels)
        acc = accuracy(outputs, labels)
        return {'loss': loss, 'acc': acc}, loss
    
    def valid_step(self, batch):
        inputs, labels = batch 
        inputs, labels = inputs.to(device), labels.to(device)
        out = self(inputs)                 
        loss = F.cross_entropy(out, labels)
        acc = accuracy(out, labels)    
        return {'val_loss': loss.detach(), 'val_acc': acc.detach()}
        

    def train_epoch_end(self, outputs):
        batch_losses = [x['loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   
        batch_accs = [x['acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()      
        return {'loss': epoch_loss.item(), 'acc': epoch_acc.item()}

    def valid_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()      
        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}

    def epoch_end(self, epoch, result, num_epochs):
        if (epoch+1) % 10 == 0 or epoch == num_epochs-1:
            print(f"Epoch [{epoch+1}] -> loss: {result['loss']:.4f}, acc: {result['acc']:.4f}, val_loss: {result['val_loss']:.4f}, val_acc: {result['val_acc']:.4f}")
    
model = MnistModel()
model.to(device)<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>def evaluate(model, val_loader):
    outputs = [model.valid_step(batch) for batch in val_loader]
    return model.valid_epoch_end(outputs)

def fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    history = []
    optimizer = opt_func(model.parameters(), lr)
    train_history = []

    for epoch in range(epochs):
        # Training Phase 
        
        for batch in train_loader:
            info, loss = model.train_step(batch)
            
            # contains batch loss and acc for training phase
            train_history.append(info)
            
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

        train_result = model.train_epoch_end(train_history)
        val_result = evaluate(model, val_loader)
        result = {**train_result, **val_result}

        model.epoch_end(epoch, result, epochs)
        history.append(result)
    return history
result = evaluate(model, val_loader)
result
history = fit(100, 0.001, model, train_loader, val_loader)
train_losses = [r['loss'] for r in history]
val_losses = [r['val_loss'] for r in history]
train_acc = [r['acc'] for r in history]
val_acc = [r['val_acc'] for r in history]

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_losses, c='k', label="train_loss")
plt.plot(val_losses, c='r', label="val_loss")
plt.legend()
plt.xlabel('epoch')
plt.ylabel('loss')
plt.title('Loss vs. epochs')

plt.subplot(1, 2, 2)
plt.plot(train_acc, c='k', label="train_acc")
plt.plot(val_acc, c='r', label="val_acc")
plt.legend()
plt.xlabel('epoch')
plt.ylabel('accuracy')
plt.title('Accuracy vs. epochs')
# Evaluate on test dataset
result = evaluate(model, test_loader)
result<jupyter_output><empty_output><jupyter_text>Prediction<jupyter_code>def predict_image(img, model):
    xb = img.unsqueeze(0)
    yb = model(xb)
    _, preds  = torch.max(yb, dim=1)
    return preds[0].item()
img, label = test_ds[919]
plt.imshow(img[0], cmap='gray')
plt.axis("off")
print('Label:', label, ', Predicted:', predict_image(img.to(device), model))<jupyter_output>Label: 2 , Predicted: 2<jupyter_text>Save Model<jupyter_code>torch.save(model.state_dict(), 'mnist-logistic.pth')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1451192623772299, 'fraction_numerical': 0.01102425335738625, 'mean_word_length': 3.557077625570776, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '590596', 'license': 'MIT', 'n_tokens_mistral': 1777, 'n_tokens_neox': 1724, 'n_words': 349, 'path': 'Week 2/mnist-logistic-minimal.ipynb', 'repo_name': 'veb-101/pytorch-tutorial', 'size': '4989'}