CIFAR_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_cifar_2-10_fixed"

train_set = datasets.CIFAR10(root=DATA_ROOT / 'data', train=True, download=True, transform=transforms.ToTensor())
test_set = datasets.CIFAR10(root=DATA_ROOT / 'data', train=False, download=True, transform=transforms.ToTensor())

# Dataloaders
train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, num_workers=6, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, num_workers=6, shuffle=False)

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