Train ResNet on CIFAR10 in a single file using PyTorch

resnet_cifar10.py

"""Train ResNet on CIFAR10 in a single file using PyTorch."""

import argparse
import json
import os
import pandas as pd
import time

import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader, random_split

import torchvision
from torchvision.transforms import Compose, ToTensor, Normalize

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"{device = }")

parser = argparse.ArgumentParser()
parser.add_argument("--num_epochs", type=int, default=11)
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--loc", type=str)
args = parser.parse_args()

# Load Data
train_split = 0.8
classes = (
    "plane",
    "car",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)
normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
fulltrainset = torchvision.datasets.CIFAR10(
    root="./data", train=True, download=True, transform=Compose([ToTensor(), normalize])
)
train_split_len = int(train_split * len(fulltrainset))
val_split_len = len(fulltrainset) - train_split_len
trainset, valset = random_split(fulltrainset, [train_split_len, val_split_len])
testset = torchvision.datasets.CIFAR10(
    root="./data",
    train=False,
    download=True,
    transform=Compose([ToTensor(), normalize]),
)
trainloader = DataLoader(
    trainset, batch_size=32, shuffle=True, num_workers=2, pin_memory=True
)
valloader = DataLoader(
    valset, batch_size=32, shuffle=True, num_workers=2, pin_memory=True
)
testloader = DataLoader(
    testset, batch_size=32, shuffle=True, num_workers=2, pin_memory=True
)

# Load Model
model = torchvision.models.resnet34()
in_ftrs = model.fc.in_features
model.fc = nn.Linear(in_ftrs, 10)
model = model.to(device)

# Train Model
save_loc = f"{args.loc}_wts.pt"
stats_loc = f"{args.loc}_stats.csv"
if os.path.exists(save_loc):
    model.load_state_dict(torch.load(save_loc))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
stats = {
    "epoch_num": [],
    "epoch_phase": [],
    "epoch_loss": [],
    "epoch_acc": [],
    "epoch_time": [],
}
for epoch in range(args.num_epochs):
    for phase in ["train", "val"]:
        if phase == "train":
            model.train()
        else:
            model.eval()
        start = time.time()
        losses = 0.0
        corrects = 0
        data_loader = trainloader if phase == "train" else valloader
        for inps, labels in data_loader:
            inps = inps.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()

            if phase == "train":
                logits = model(inps)
            else:
                with torch.no_grad():
                    logits = model(inps)
            _, preds = torch.max(logits, 1)
            loss = criterion(logits, labels)

            if phase == "train":
                loss.backward()
                optimizer.step()

            losses += loss.item() * inps.size(0)
            corrects += torch.sum(preds == labels.data)
        data_len = len(trainset) if phase == "train" else len(valset)
        epoch_loss = losses / data_len
        epoch_acc = corrects.double() / data_len
        epoch_time = time.time() - start
        print(
            f"{epoch}/{args.num_epochs} {phase}: {epoch_loss=:.4f}, {epoch_acc=:.4f}, {epoch_time=:.4f}"
        )
        stats["epoch_num"].append(epoch)
        stats["epoch_phase"].append(phase)
        stats["epoch_loss"].append(epoch_loss)
        stats["epoch_acc"].append(epoch_acc.item())
        stats["epoch_time"].append(epoch_time)

    if not epoch % 5:
        torch.save(model.state_dict(), save_loc)
        print(f"Saved model parameters to {save_loc}")
        pd.DataFrame(stats).to_csv(stats_loc, index=False)
        print(f"Saved training statistics to {stats_loc}")

# Evaluate Model
model.eval()
losses = 0.0
corrects = 0
start = time.time()
test_stats_loc = f"{args.loc}_test_stats.json"
for inps, labels in testloader:
    inps = inps.to(device)
    labels = labels.to(device)

    with torch.no_grad():
        logits = model(inps)
    _, preds = torch.max(logits, 1)
    loss = criterion(logits, labels)

    losses += loss.item() * inps.size(0)
    corrects += torch.sum(preds == labels.data)
test_loss = losses / len(testset)
test_acc = corrects.double() / len(testset)
test_time = time.time() - start
print(f"Test Stats: {test_loss=:.4f}, {test_acc=:.4f}, {test_time=:.4f}")
with open(test_stats_loc, "w") as fh:
    json.dump(
        {"test_loss": test_loss, "test_acc": test_acc.item(), "test_time": test_time},
        fh,
    )
print(f"Saved test statistics to {test_stats_loc}")

resnet_cifar10_script.py

"""Script to run parallel slurm jobs"""

import argparse
import submitit
from pathlib import Path

parser = argparse.ArgumentParser(description="ResNet CIFAR-10")
parser.add_argument("--num_epochs", type=int, default=11)
parser.add_argument("--job_name", type=str, default="job_%j")
parser.add_argument("--pathname", type=str, default="results")

args = parser.parse_args()

Path(args.pathname).mkdir(parents=True, exist_ok=True)

executor = submitit.SlurmExecutor(folder=args.job_name)
executor.update_parameters(
    time=4000,
    cpus_per_task=2,
    gres="gpu:1",
    job_name="resnet_cifar10",
    ntasks_per_node=1,
    exclude="locus-1-[21,25,29]",
    mem="12G",
    array_parallelism=4,
)
with executor.batch():
    lrs = [1e-2, 1e-3, 1e-4, 1e-5, 1e-6]
    fnames = ["tmp2", "tmp3", "tmp4", "tmp5", "tmp6"]
    for lr, loc in zip(lrs, fnames):
        function = submitit.helpers.CommandFunction(
            [
                "python",
                "-u",
                "resnet_cifar10.py",
                "--num_epochs",
                f"{args.num_epochs}",
                "--lr",
                f"{lr}",
                "--loc",
                f"{args.pathname}/{loc}",
            ]
        )
        print(f"Submitting job: {' '.join(function.command)}")
        job = executor.submit(function)

stats.csv

test_stats.json

{"test_loss": 2.019672994709015, "test_acc": 0.6923, "test_time": 1.195981502532959}