add simple neural network and ResNet for Pytorch

Pytorch_MNIST.py

@@ -0,0 +1,158 @@
+from __future__ import print_function
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.optim.lr_scheduler import StepLR
+import matplotlib.pyplot as plt
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = F.log_softmax(x, dim=1)
+        return output
+
+train_losses = []
+train_counter = []
+test_losses = []
+
+def train(args, 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 = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+            train_losses.append(loss.item())
+            train_counter.append((batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))
+            if args.dry_run:
+                break
+
+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 += F.nll_loss(output, target, reduction='sum').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)
+    test_losses.append(test_loss)
+
+    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        100. * correct / len(test_loader.dataset)))
+
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
+                        help='input batch size for testing (default: 1000)')
+    parser.add_argument('--epochs', type=int, default=3, metavar='N',
+                        help='number of epochs to train (default: 14)')
+    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
+                        help='learning rate (default: 1.0)')
+    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
+                        help='Learning rate step gamma (default: 0.7)')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--no-mps', action='store_true', default=False,
+                        help='disables macOS GPU training')
+    parser.add_argument('--dry-run', action='store_true', default=False,
+                        help='quickly check a single pass')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+    parser.add_argument('--save-model', action='store_true', default=False,
+                        help='For Saving the current Model')
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+    use_mps = not args.no_mps and torch.backends.mps.is_available()
+
+    torch.manual_seed(args.seed)
+
+    if use_cuda:
+        device = torch.device("cuda")
+    elif use_mps:
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    train_kwargs = {'batch_size': args.batch_size}
+    test_kwargs = {'batch_size': args.test_batch_size}
+    if use_cuda:
+        cuda_kwargs = {'num_workers': 1,
+                       'pin_memory': True,
+                       'shuffle': True}
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    transform=transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,))
+        ])
+    dataset1 = datasets.MNIST('../data', train=True, download=True,
+                       transform=transform)
+    dataset2 = datasets.MNIST('../data', train=False,
+                       transform=transform)
+    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
+    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+    model = Net().to(device)
+    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+        test(model, device, test_loader)
+        scheduler.step()
+
+    if args.save_model:
+        torch.save(model.state_dict(), "mnist_cnn.pt")
+    fig = plt.figure()
+    test_counter = [i*len(train_loader.dataset) for i in range(epoch + 1)]
+    plt.plot(train_counter, train_losses, color='blue')
+    plt.scatter(test_counter, test_losses, color='red')
+    plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
+    plt.xlabel('number of training examples seen')
+    plt.ylabel('negative log likelihood loss')
+
+
+if __name__ == '__main__':
+    main()

Pytorch_ResNet.py

@@ -0,0 +1,239 @@
+import os
+import time
+
+import numpy as np
+import pandas as pd
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from torchvision import datasets
+from torchvision import transforms
+
+import matplotlib.pyplot as plt
+from PIL import Image
+
+
+
+if torch.cuda.is_available():
+    torch.backends.cudnn.deterministic = True
+
+##########################
+### SETTINGS
+##########################
+
+# Hyperparameters
+RANDOM_SEED = 1
+LEARNING_RATE = 0.001
+BATCH_SIZE = 128
+NUM_EPOCHS = 10
+
+# Architecture
+NUM_FEATURES = 28*28
+NUM_CLASSES = 10
+
+# Other
+DEVICE = "cuda:1" if torch.cuda.is_available() else "cpu"
+GRAYSCALE = True
+
+##########################
+### MNIST DATASET
+##########################
+
+# Note transforms.ToTensor() scales input images
+# to 0-1 range
+train_dataset = datasets.MNIST(root='data', 
+                               train=True, 
+                               transform=transforms.ToTensor(),
+                               download=True)
+
+test_dataset = datasets.MNIST(root='data', 
+                              train=False, 
+                              transform=transforms.ToTensor())
+
+
+train_loader = DataLoader(dataset=train_dataset, 
+                          batch_size=BATCH_SIZE, 
+                          shuffle=True)
+
+test_loader = DataLoader(dataset=test_dataset, 
+                         batch_size=BATCH_SIZE, 
+                         shuffle=False)
+
+##########################
+### MODEL
+##########################
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes, grayscale):
+        self.inplanes = 64
+        if grayscale:
+            in_dim = 1
+        else:
+            in_dim = 3
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, (2. / n)**.5)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    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)
+        # because MNIST is already 1x1 here:
+        # disable avg pooling
+        #x = self.avgpool(x)
+
+        x = x.view(x.size(0), -1)
+        logits = self.fc(x)
+        probas = F.softmax(logits, dim=1)
+        return logits, probas
+
+def resnet18(num_classes):
+    """Constructs a ResNet-18 model."""
+    model = ResNet(block=BasicBlock, 
+                   layers=[2, 2, 2, 2],
+                   num_classes=NUM_CLASSES,
+                   grayscale=GRAYSCALE)
+    return model
+
+torch.manual_seed(RANDOM_SEED)
+
+model = resnet18(NUM_CLASSES)
+model.to(DEVICE)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)  
+
+
+def compute_accuracy(model, data_loader, device):
+    correct_pred, num_examples = 0, 0
+    for i, (features, targets) in enumerate(data_loader):
+
+        features = features.to(device)
+        targets = targets.to(device)
+
+        logits, probas = model(features)
+        _, predicted_labels = torch.max(probas, 1)
+        num_examples += targets.size(0)
+        correct_pred += (predicted_labels == targets).sum()
+    return correct_pred.float()/num_examples * 100
+
+
+start_time = time.time()
+for epoch in range(NUM_EPOCHS):
+
+    model.train()
+    for batch_idx, (features, targets) in enumerate(train_loader):
+
+        features = features.to(DEVICE)
+        targets = targets.to(DEVICE)
+
+        ### FORWARD AND BACK PROP
+        logits, probas = model(features)
+        cost = F.cross_entropy(logits, targets)
+        optimizer.zero_grad()
+
+        cost.backward()
+
+        ### UPDATE MODEL PARAMETERS
+        optimizer.step()
+
+        ### LOGGING
+        if not batch_idx % 50:
+            print ('Epoch: %03d/%03d | Batch %04d/%04d | Cost: %.4f' 
+                   %(epoch+1, NUM_EPOCHS, batch_idx, 
+                     len(train_loader), cost))
+
+
+
+    model.eval()
+    with torch.set_grad_enabled(False): # save memory during inference
+        print('Epoch: %03d/%03d | Train: %.3f%%' % (
+              epoch+1, NUM_EPOCHS, 
+              compute_accuracy(model, train_loader, device=DEVICE)))
+
+    print('Time elapsed: %.2f min' % ((time.time() - start_time)/60))
+
+print('Total Training Time: %.2f min' % ((time.time() - start_time)/60))
+
+with torch.set_grad_enabled(False): # save memory during inference
+    print('Test accuracy: %.2f%%' % (compute_accuracy(model, test_loader, device=DEVICE)))