demo batch norm

for train and test

# Importing necessary libraries
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define the neural network with Batch Normalization
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.fc1 = nn.Linear(64*7*7, 128)
        self.bn3 = nn.BatchNorm1d(128)
        self.fc2 = nn.Linear(128, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = torch.relu(x)
        x = torch.max_pool2d(x, 2)
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = torch.relu(x)
        x = torch.max_pool2d(x, 2)
        
        x = x.view(-1, 64*7*7)
        x = self.fc1(x)
        x = self.bn3(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

# Load the dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('/data/solai/kmgit/dwDbs/mnist', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('/data/solai/kmgit/dwDbs/mnist', train=False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# Instantiate the model, define the loss function and the optimizer
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

# Test loop
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()  # Sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # Get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Train and test the model
for epoch in range(1, 6):
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader)

Train Epoch: 1 [0/60000 (0%)]	Loss: 2.356473
Train Epoch: 1 [6400/60000 (11%)]	Loss: 0.153310
Train Epoch: 1 [12800/60000 (21%)]	Loss: 0.082770
Train Epoch: 1 [19200/60000 (32%)]	Loss: 0.020425
Train Epoch: 1 [25600/60000 (43%)]	Loss: 0.047038
Train Epoch: 1 [32000/60000 (53%)]	Loss: 0.051890
Train Epoch: 1 [38400/60000 (64%)]	Loss: 0.058858
Train Epoch: 1 [44800/60000 (75%)]	Loss: 0.111991
Train Epoch: 1 [51200/60000 (85%)]	Loss: 0.046442
Train Epoch: 1 [57600/60000 (96%)]	Loss: 0.080021

Test set: Average loss: 0.0000, Accuracy: 9882/10000 (99%)

Train Epoch: 2 [0/60000 (0%)]	Loss: 0.027627
Train Epoch: 2 [6400/60000 (11%)]	Loss: 0.024146
Train Epoch: 2 [12800/60000 (21%)]	Loss: 0.013543
Train Epoch: 2 [19200/60000 (32%)]	Loss: 0.038260
Train Epoch: 2 [25600/60000 (43%)]	Loss: 0.001468
Train Epoch: 2 [32000/60000 (53%)]	Loss: 0.070194
Train Epoch: 2 [38400/60000 (64%)]	Loss: 0.094309
Train Epoch: 2 [44800/60000 (75%)]	Loss: 0.016474
Train Epoch: 2 [51200/60000 (85%)]	Loss: 0.004236
Train Epoch: 2 [57600/60000 (96%)]	Loss: 0.041856

Test set: Average loss: 0.0000, Accuracy: 9887/10000 (99%)

Train Epoch: 3 [0/60000 (0%)]	Loss: 0.048179
Train Epoch: 3 [6400/60000 (11%)]	Loss: 0.028058
Train Epoch: 3 [12800/60000 (21%)]	Loss: 0.006302
Train Epoch: 3 [19200/60000 (32%)]	Loss: 0.002997
Train Epoch: 3 [25600/60000 (43%)]	Loss: 0.036016
Train Epoch: 3 [32000/60000 (53%)]	Loss: 0.004273
Train Epoch: 3 [38400/60000 (64%)]	Loss: 0.006739
Train Epoch: 3 [44800/60000 (75%)]	Loss: 0.001813
Train Epoch: 3 [51200/60000 (85%)]	Loss: 0.004095
Train Epoch: 3 [57600/60000 (96%)]	Loss: 0.010136

Test set: Average loss: 0.0000, Accuracy: 9913/10000 (99%)

Train Epoch: 4 [0/60000 (0%)]	Loss: 0.004269
Train Epoch: 4 [6400/60000 (11%)]	Loss: 0.004508
Train Epoch: 4 [12800/60000 (21%)]	Loss: 0.002657
Train Epoch: 4 [19200/60000 (32%)]	Loss: 0.008510
Train Epoch: 4 [25600/60000 (43%)]	Loss: 0.109596
Train Epoch: 4 [32000/60000 (53%)]	Loss: 0.010197
Train Epoch: 4 [38400/60000 (64%)]	Loss: 0.016853
Train Epoch: 4 [44800/60000 (75%)]	Loss: 0.004828
Train Epoch: 4 [51200/60000 (85%)]	Loss: 0.000939
Train Epoch: 4 [57600/60000 (96%)]	Loss: 0.003024

Test set: Average loss: 0.0000, Accuracy: 9890/10000 (99%)

Train Epoch: 5 [0/60000 (0%)]	Loss: 0.003337
Train Epoch: 5 [6400/60000 (11%)]	Loss: 0.009937
Train Epoch: 5 [12800/60000 (21%)]	Loss: 0.002990
Train Epoch: 5 [19200/60000 (32%)]	Loss: 0.007060
Train Epoch: 5 [25600/60000 (43%)]	Loss: 0.007873
Train Epoch: 5 [32000/60000 (53%)]	Loss: 0.000431
Train Epoch: 5 [38400/60000 (64%)]	Loss: 0.003605
Train Epoch: 5 [44800/60000 (75%)]	Loss: 0.002747
Train Epoch: 5 [51200/60000 (85%)]	Loss: 0.005132
Train Epoch: 5 [57600/60000 (96%)]	Loss: 0.005143

Test set: Average loss: 0.0000, Accuracy: 9916/10000 (99%)