demo_mnist.py

from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets, transforms


def get_device() -> torch.device:
    device_arg = "cpu"
    if torch.backends.mps.is_available():
        device_arg = "mps"
        print("MPS device detected")

        macos_version = torch.backends.mps.is_macos13_or_newer()
        print(f"macOS 13 or newer: {macos_version}")

    elif torch.cuda.is_available():
        device_arg = "cuda"
        print("CUDA device detected")

        gpu_count = torch.cuda.device_count()
        print(f"Number of available GPUs: {gpu_count}")

        for i in range(gpu_count):
            torch.cuda.set_device(i)

            props = torch.cuda.get_device_properties(i)

            print(f"\nGPU {i}:")
            print(f"  Name: {props.name}")
            print(f"  Compute Capability: {props.major}.{props.minor}")
            print(f"  Total Memory: {props.total_memory / 1024**3:.2f} GB")
            print(f"  Multi-Processor Count: {props.multi_processor_count}")

            memory_allocated = torch.cuda.memory_allocated(i) / 1024**2
            memory_reserved = torch.cuda.memory_reserved(i) / 1024**2
            print(f"  Memory Allocated: {memory_allocated:.2f} MB")
            print(f"  Memory Reserved: {memory_reserved:.2f} MB")

            utilization = torch.cuda.utilization(i)
            print(f"  GPU Utilization: {utilization}%")

    device = torch.device(device_arg)
    return device


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.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(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.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


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.0 * batch_idx / len(train_loader), loss.item()))


def test(args, 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)
            # sum up batch loss
            test_loss += F.nll_loss(output, target, reduction="sum").item()
            # get the index of the max log-probability
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

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


def main():
    args = {
        "batch_size": 64,
        "test_batch_size": 1000,
        "epochs": 2,
        "lr": 1.0,
        "gamma": 0.7,
        "seed": 1,
        "log_interval": 25,
    }

    torch.manual_seed(args["seed"])

    kwargs = {"num_workers": 1, "pin_memory": True}
    train_loader = torch.utils.data.DataLoader(datasets.MNIST("./data", train=True, download=True, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])), batch_size=args["batch_size"], shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(datasets.MNIST("./data", train=False, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])), batch_size=args["test_batch_size"], shuffle=True, **kwargs)

    device = get_device()
    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(args, model, device, test_loader)
        scheduler.step()


if __name__ == "__main__":
    main()