Example 3: PyTorch + MNIST + GPU

This example demonstrates the recommended workflow for running a GPU deep learning job on the CoSTAR cluster: using Apptainer containers.

The following example trains a convolutional neural network on the MNIST dataset (handwritten digits). It requests a GPU and runs a Python script inside a public PyTorch container from Docker Hub. This avoids the need to maintain a Conda installation on the cluster or build your own Apptainer or Docker image with a Conda environment and PyTorch packages.

1. Connect to the CoSTAR cluster

First, open your terminal and SSH into the CoSTAR login node:

ssh ab1234@costar-login01

Replace ab1234 with your username.

2. Create the submission script

Download the example submission script, or create a file called example-03-gpu.sh in your working directory and copy in the script below.

Notice the apptainer exec command at the bottom of the submission script. The docker://pytorch/pytorch:latest image reference tells Apptainer to pull the public PyTorch container from Docker Hub and run the Python script inside that container.

Note

This example uses the public pytorch/pytorch:latest Docker image from Docker Hub. Apptainer automatically converts the Docker image into an Apptainer-compatible image and caches it in the standard Apptainer cache in your home directory.

The first run may take longer while Apptainer pulls and converts the image. Later runs are usually faster because the converted image can be reused from the cache.

example-03-gpu.sh (click to collapse / view)
 1#!/bin/sh
 2
 3#SBATCH --job-name="GPU PyTorch"
 4#SBATCH --partition=main
 5#SBATCH --gpus=1
 6#SBATCH --nodes=1
 7#SBATCH --ntasks-per-node=1
 8#SBATCH --mem=20G
 9#SBATCH --time=00-00:25:00
10#SBATCH -o slurm.%N.%j.out
11#SBATCH -e slurm.%N.%j.err
12
13apptainer exec docker://pytorch/pytorch:latest python example-03-gpu.py

3. The Python script

This is a standard PyTorch training script. Download the example Python script, or create a file called example-03-gpu.py in your working directory and copy in the script below.

example-03-gpu.py (click to collapse / view)
  1from __future__ import print_function
  2import argparse
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6import torch.optim as optim
  7from torchvision import datasets, transforms
  8
  9# Training settings
 10parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
 11parser.add_argument('--batch-size', type=int, default=64, metavar='N',
 12                    help='input batch size for training (default: 64)')
 13parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
 14                    help='input batch size for testing (default: 1000)')
 15parser.add_argument('--epochs', type=int, default=10, metavar='N',
 16                    help='number of epochs to train (default: 10)')
 17parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
 18                    help='learning rate (default: 0.01)')
 19parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
 20                    help='SGD momentum (default: 0.5)')
 21parser.add_argument('--seed', type=int, default=1, metavar='S',
 22                    help='random seed (default: 1)')
 23parser.add_argument('--log-interval', type=int, default=10, metavar='N',
 24                    help='how many batches to wait before logging training status')
 25args = parser.parse_args()
 26
 27if not torch.cuda.is_available():
 28    raise RuntimeError("CUDA is not available. Check that the job requested a GPU.")
 29
 30device = torch.device("cuda")
 31print(f"CUDA available: {torch.cuda.is_available()}")
 32print(f"CUDA device count: {torch.cuda.device_count()}")
 33print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
 34
 35torch.manual_seed(args.seed)
 36
 37train_loader = torch.utils.data.DataLoader(
 38    datasets.MNIST('data', train=True, download=True,
 39                   transform=transforms.Compose([
 40                       transforms.ToTensor(),
 41                       transforms.Normalize((0.1307,), (0.3081,))
 42                   ])),
 43    batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True)
 44test_loader = torch.utils.data.DataLoader(
 45    datasets.MNIST('data', train=False, transform=transforms.Compose([
 46                       transforms.ToTensor(),
 47                       transforms.Normalize((0.1307,), (0.3081,))
 48                   ])),
 49    batch_size=args.test_batch_size, shuffle=True, num_workers=2, pin_memory=True)
 50
 51
 52class Net(nn.Module):
 53    def __init__(self):
 54        super(Net, self).__init__()
 55        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
 56        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
 57        self.conv2_drop = nn.Dropout2d()
 58        self.fc1 = nn.Linear(320, 50)
 59        self.fc2 = nn.Linear(50, 10)
 60
 61    def forward(self, x):
 62        x = F.relu(F.max_pool2d(self.conv1(x), 2))
 63        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
 64        x = x.view(-1, 320)
 65        x = F.relu(self.fc1(x))
 66        x = F.dropout(x, training=self.training)
 67        x = self.fc2(x)
 68        return F.log_softmax(x, dim=1)
 69
 70model = Net().to(device)
 71
 72optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
 73
 74def train(epoch):
 75    model.train()
 76    for batch_idx, (data, target) in enumerate(train_loader):
 77        data, target = data.to(device), target.to(device)
 78        optimizer.zero_grad()
 79        output = model(data)
 80        loss = F.nll_loss(output, target)
 81        loss.backward()
 82        optimizer.step()
 83        if batch_idx % args.log_interval == 0:
 84            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
 85                epoch, batch_idx * len(data), len(train_loader.dataset),
 86                100. * batch_idx / len(train_loader), loss.item()))
 87
 88def test():
 89    model.eval()
 90    test_loss = 0
 91    correct = 0
 92    with torch.no_grad():
 93        for data, target in test_loader:
 94            data, target = data.to(device), target.to(device)
 95            output = model(data)
 96            test_loss += F.nll_loss(output, target, size_average=False).item()
 97            pred = output.max(1, keepdim=True)[1]
 98            correct += pred.eq(target.view_as(pred)).sum().item()
 99
100    test_loss /= len(test_loader.dataset)
101    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
102        test_loss, correct, len(test_loader.dataset),
103        100. * correct / len(test_loader.dataset)))
104
105
106for epoch in range(1, args.epochs + 1):
107    train(epoch)
108    test()

4. Submit your job

Submit your script to the Slurm scheduler using the following command:

sbatch example-03-gpu.sh

You can check the status of your job by running:

squeue --me

5. The output

Once the job completes, you will see the Slurm output and error files in your directory:

  • slurm.<node>.<job_id>.out / .err

    These are the standard Slurm logs. The Python script’s training progress, including epoch numbers, loss, and accuracy, is written to the .out file. The .out file should also show that CUDA was detected and name the GPU device used by PyTorch. If the container failed to pull, or the scheduler encountered an issue, the error will be detailed in the .err file.

Note

If the job is using a GPU and CUDA is available, the output file should include lines similar to:

CUDA available: True
CUDA device count: 1
Using CUDA device: NVIDIA H200
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.327496
Train Epoch: 1 [640/60000 (1%)] Loss: 2.328413