Skip to content
Snippets Groups Projects
main.py 2.12 KiB
Newer Older
import dataloader
import model
import test
from train import train
import yaml
import losses
Yandi's avatar
Yandi committed
import optimizers
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

def optimizer(cfg, model):
    result = {"Adam" : torch.optim.Adam(model.parameters())}
    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
    if use_cuda :
        device = torch.device('cuda')
    else :
        device = toch.device('cpu')

Yandi's avatar
Yandi committed
    #model = model.build_model(cfg, 18)

    model = nn.Sequential(
        nn.Linear(18,1,False),
        nn.ReLU()
    )
Yandi's avatar
Yandi committed
    model = model.to(device)
Yandi's avatar
Yandi committed
    for param in list(model.parameters()):
        param = 1

Yandi's avatar
Yandi committed
    f_loss = losses.RMSLELoss()
Yandi's avatar
Yandi committed
    optimizer = optimizer(cfg, model)
Yandi's avatar
Yandi committed
    #optimizer = torch.optim.Adam((model.parameters()), lr = 10000)
Yandi's avatar
Yandi committed
  
Yandi's avatar
Yandi committed
    for t in range(cfg["Training"]["Epochs"]):
        print("Epoch {}".format(t))
        train(model, train_loader, f_loss, optimizer, device)
Yandi's avatar
Yandi committed

        print(list(model.parameters())[0].grad)
Yandi's avatar
Yandi committed
        val_loss = test.test(model, valid_loader, f_loss, device)
        print(" Validation : Loss : {:.4f}".format(val_loss))
    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)
    """