MNIST_ResNet.py

import torch
from torchvision import datasets, transforms

from modules.parser import DATA_ROOT
from modules.resnet import resnet18

CHECKPOINTS_FOLDER = DATA_ROOT / "checkpoints/ResNet_mnist_fixed"

transform = transforms.Compose(
    [transforms.Grayscale(num_output_channels=3), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ])

batch_size = 128
train_set = datasets.MNIST(root=DATA_ROOT / 'data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2)

test_set = datasets.MNIST(root=DATA_ROOT / 'data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2)

NAME = "model"
EXTENSION = ".pt"
Model = resnet18(pretrained=True)