lenet_mnist / main.py
rzimmerdev's picture
Final product
46627d2
raw
history blame contribute delete
752 Bytes
import torch
from src.models import CNN
from src.dataset import DatasetMNIST, download_mnist
from src.train import get_dataloaders, train_net_manually, train_net_lightning
def main(device):
mnist = download_mnist("downloads/mnist/")
dataset, test_data = DatasetMNIST(*mnist["train"]), DatasetMNIST(*mnist["test"])
train_loader, validate_loader, test_loader = get_dataloaders(dataset, test_data)
# Training manually
net = CNN(input_channels=1, num_classes=10).to(device)
optim = torch.optim.Adam(net.parameters(), lr=1e-4)
loss_fn = torch.nn.CrossEntropyLoss()
max_epochs = 1
train_net_manually(net, optim, loss_fn, train_loader, validate_loader, max_epochs, device)
if __name__ == "__main__":
main("cpu")