Skip to content
Snippets Groups Projects
main.py 3.02 KiB
Newer Older
Yandi's avatar
Yandi committed
#Internal imports
import dataloader
import model
import test
from train import train
import losses
Yandi's avatar
Yandi committed
import optimizers
Yandi's avatar
Yandi committed
import create_submission
import utils

#External imports
import yaml
Yandi's avatar
Yandi committed
import torch
import logging
Yandi's avatar
Yandi committed
import torch.optim
Yandi's avatar
Yandi committed
import torch.nn as nn
Yandi's avatar
Yandi committed
import os
Yandi's avatar
Yandi committed

Yandi's avatar
Yandi committed
def optimizer(cfg, network):
    result = {"Adam" : torch.optim.Adam(network.parameters())}
Yandi's avatar
Yandi committed
    return result[cfg["Optimizer"]]

if __name__ == "__main__":
Yandi's avatar
Yandi committed
    logging.basicConfig(filename='logs/main_unit_test.log', level=logging.INFO)

    config_file = open("config.yml")
    cfg = yaml.load(config_file)

    use_cuda = torch.cuda.is_available()
    trainpath           = cfg["Dataset"]["_DEFAULT_TRAIN_FILEPATH"]
Yandi's avatar
Yandi committed
    num_days            = int(cfg["Dataset"]["num_days"])
    batch_size          = int(cfg["Dataset"]["batch_size"])
    num_workers         = int(cfg["Dataset"]["num_workers"])
    valid_ratio         = float(cfg["Dataset"]["valid_ratio"])
    max_num_samples     = eval(cfg["Dataset"]["max_num_samples"])

    train_loader, valid_loader = dataloader.get_dataloaders(
        trainpath,
        num_days,
        batch_size,
        num_workers,
        use_cuda,
        valid_ratio,
        overwrite_index=True,
Yandi's avatar
Yandi committed
        max_num_samples=max_num_samples,
        train_transform=dataloader.transform_remove_space_time(),
        valid_transform=dataloader.transform_remove_space_time()
    if use_cuda :
        device = torch.device('cuda')
    else :
        device = toch.device('cpu')

Yandi's avatar
Yandi committed
    #network = network.build_network(cfg, 18)
Yandi's avatar
Yandi committed
    network = nn.Sequential(
Yandi's avatar
Yandi committed
        nn.Linear(14,8,False),
        nn.ReLU(),
Yandi's avatar
Yandi committed
        nn.Linear(8, 35, True),
        nn.ReLU(),
        nn.Linear(35,35,True),
        nn.ReLU(),
        nn.Linear(35,35,True),
        nn.ReLU(),
        nn.Linear(35,35,True),
Yandi's avatar
Yandi committed
        nn.ReLU(),
Yandi's avatar
Yandi committed
        nn.Linear(35,35,True),
Yandi's avatar
Yandi committed
        nn.ReLU(),
Yandi's avatar
Yandi committed
        nn.Linear(35,1, True),
        nn.ReLU()
Yandi's avatar
Yandi committed
    )
Yandi's avatar
Yandi committed
    def init_xavier(module):
        if type(module)==nn.Linear:
            nn.init.xavier_uniform_(module.weight)
    network = network.to(device)

    """
    for param in list(network.parameters()):
Yandi's avatar
Yandi committed
        param = 1
Yandi's avatar
Yandi committed
    """
Yandi's avatar
Yandi committed
    f_loss = losses.RMSLELoss()
Yandi's avatar
Yandi committed
    optimizer = optimizer(cfg, network)

    logdir = utils.create_unique_logpath(cfg["LogDir"], cfg["Model"]["Name"])
    network_checkpoint = model.ModelCheckpoint(logdir + "/best_model.pt", network)
Yandi's avatar
Yandi committed
    for t in range(cfg["Training"]["Epochs"]):
Yandi's avatar
Yandi committed
        torch.autograd.set_detect_anomaly(True)
Yandi's avatar
Yandi committed
        print("Epoch {}".format(t))
Yandi's avatar
Yandi committed
        train(network, train_loader, f_loss, optimizer, device)
Yandi's avatar
Yandi committed
        #print(list(network.parameters())[0].grad)
        val_loss = test.test(network, valid_loader, f_loss, device)

        network_checkpoint.update(val_loss)

Yandi's avatar
Yandi committed
        print(" Validation : Loss : {:.4f}".format(val_loss))
Yandi's avatar
Yandi committed

    create_submission.create_submission(network, None)
    logdir = generate_unique_logpath(top_logdir, "linear")
    print("Logging to {}".format(logdir))
    # -> Prints out     Logging to   ./logs/linear_1
    if not os.path.exists(logdir):
        os.mkdir(logdir)
    """