main.py

import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

from activation import load_actF
import models
from utils import Tag, Logger

epochs = 10
batch_size = 64
valid_split = 0.2
num_workers = 4
pin_memory = True
test_batch_size = 1000
topk = (1,5)
verbose = True
model_name = 'Net' #'ResNet18'
activation_function = 'relu'
afkw = dict(
    # a=0.2
    )

def train(dataset, device, net, criterion, optimizer, valid_split=valid_split):
    N, r = len(dataset), valid_split
    trainset, validset = random_split(dataset, [N-int(N*r), int(N*r)])
    train_loader = DataLoader(dataset=trainset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory)
    valid_loader = DataLoader(dataset=validset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory)

    tag = Tag(topk=topk, verbose=verbose)
    net.train()
    for inputs, targets in tag(train_loader):
        inputs.to(device), targets.to(device)
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_log = tag.log(outputs, targets, loss)
    
    net.eval()
    with torch.no_grad():
        for inputs, targets in tag(valid_loader):
            inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            valid_log = tag.log(outputs, targets, loss)
    
    return train_log, valid_log

def test(evalset, device, net, criterion):
    test_loader = DataLoader(dataset=evalset, batch_size=test_batch_size, num_workers=num_workers, pin_memory=pin_memory)
    
    tag = Tag(topk=topk, verbose=verbose)
    net.eval()
    with torch.no_grad():
        for inputs, targets in tag(test_loader):
            inputs.to(device), targets.to(device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            test_log = tag.log(outputs, targets, loss)
    
    return test_log


def main():
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.CIFAR10(root='../data', train=True, transform= transform, download=True)
    evalset = datasets.CIFAR10(root='../data', train=False, transform= transform, download=True)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = models.Net()
    net.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters())

    logger = Logger()
    for epoch in logger(epochs):
        train_log, valid_log = train(dataset, device, net, criterion, optimizer)
        logger.log_etv(epoch, train_log, valid_log)

    test_log = test(evalset, device, net, criterion)
    logger.log_tst(test_log)

if __name__ == "__main__":
    main()