<reponame>Avonite/context-project<filename>week3/train.py
import numpy as np
import torch
import torchvision
from sklearn.model_selection import train_test_split
from torch.utils.data import Subset, DataLoader
from torchvision import transforms

from week4.network import NeuralNetwork
from week4.settings import EPOCHS
from week4.util import train_and_test

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

dataset = torchvision.datasets.FashionMNIST('./data/', download=True, train=True, transform=transform)

# Split the indices in a stratified way
indices = np.arange(len(dataset))
train_indices, test_indices = train_test_split(indices, train_size=100*10, stratify=dataset.targets)

# Wrap into Subsets and DataLoaders
train_dataset = Subset(dataset, train_indices)
test_dataset = Subset(dataset, test_indices)

train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=10)
test_dataloader = DataLoader(train_dataset, shuffle=False, batch_size=10)

train_and_test(NeuralNetwork(), train_dataloader, test_dataloader, epochs=EPOCHS)
