train-gpu.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 22 00:07:48 2022

@author: alitaghibakhshi
"""
import sys
sys.path.insert(1, 'utils/')
import matplotlib.pyplot as plt
import numpy as np
import torch
from pathlib import Path
import os
import os.path
import torch as T
import copy
import random
import scipy
from grids_gpu import *
from utils_gpu import *
import argparse
from mggnn_gpu import *
from lloyd_gunet import *
import time

manual_seeding = 34210
np.random.seed(manual_seeding)
random.seed(manual_seeding)
torch.manual_seed(manual_seeding)
# torch.use_deterministic_algorithms(True)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False


if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')
    
    
train_parser = argparse.ArgumentParser(description='Settings for training machine learning for ORAS')

train_parser.add_argument('--num-epoch', type=int, default=10, help='Number of training epochs')
train_parser.add_argument('--mini-batch-size', type=int, default=1, help='Coarsening ratio for aggregation')
train_parser.add_argument('--lr', type=float, default= 5e-4, help='Learning rate')
train_parser.add_argument('--TAGConv-k', type=int, default=2, help='TAGConv # of hops')
train_parser.add_argument('--dim', type=int, default=128, help='Dimension of TAGConv filter')
train_parser.add_argument('--data-set', type=str, default='Data/old_data', help='Directory of the training data')
train_parser.add_argument('--K', type=int, default=10, help='Number of iterations in the loss function')
train_parser.add_argument('--GNN', type=str, default='MG-GNN', help='MG-GNN or Graph-Unet')
train_parser.add_argument('--path', type=str, default='Models/old-model', help='MG-GNN or Graph-Unet')
train_parser.add_argument('--lvl', type=int, default=2, help='Number of Levels')
train_parser.add_argument('--layer', type=int, default=4, help='Number of Layers')

train_args = train_parser.parse_args()

if __name__ == "__main__":
    
        
    path = train_args.path
    
    if not os.path.exists(path):
        os.makedirs(path)
    
    list_grids = []
    
    num_data = 1#sum((len(f) for _, _, f in os.walk(train_args.data_set)))-1

    for i in range(num_data):

        g = torch.load(train_args.data_set+"/grid"+str(i)+".pth")
        list_grids.append(g)

    print('Finished Uploading Training Data')
    
    if train_args.GNN == 'MG-GNN':
        model = MGGNN(lvl=train_args.lvl, dim_embed=128, num_layers=train_args.layer, K=train_args.TAGConv_k, ratio=0.2, lr=train_args.lr)
    elif train_args.GNN == 'Graph-Unet':
        model = lloyd_gunet(train_args.lvl, train_args.layer, 128, K = 2, ratio = 0.2, lr = train_args.lr)
    else:
        raise ValueError("Select GNN architecture between MG-GNN and Graph-Unet")
    
    model.load_state_dict(torch.load('Models/model_epoch_best.pth'), strict=False)

    model.to(device)

    print('Number of parameters: ',sum(p.numel() for p in model.parameters()))

    epoch_loss_list = []
    all_indices = np.arange(num_data)
    model.optimizer.zero_grad()
    
    current_best_loss = 10**12


    epoch_loss = 0

    for epoch in range(train_args.num_epoch):

        loss = 0
        np.random.shuffle(all_indices)
        mbs = train_args.mini_batch_size
        print("Epoch = ", epoch)
        print("-----------------")
        for count in range(int(np.ceil((num_data)/mbs))):

            batch_idxs = all_indices[count*mbs:min((count+1)*mbs, num_data)]
            t1 = time.time()
            for i in batch_idxs:

                grid = list_grids[i]
                grid.to(device)
                output = model.forward(grid, train = True)
                sump = 0
                for p in model.parameters():
                    sump+=abs(p.flatten()).sum()
                # print(np.random.rand())
                # print(torch.rand(3))
                # print(sump)
                # print(abs(grid.A).toarray().flatten().sum())
                # print(output[0].to_dense().flatten().sum())
                # print(output[1].to_dense().flatten().sum())
                # sys.exit()
                u = torch.rand(grid.x.shape[0],100).double().to(device)
                u = u/(((u**2).sum(0))**0.5).unsqueeze(0)

                current_loss = stationary_max(grid, output, u = u, K = train_args.K, precond_type='ML_ORAS')

                loss += current_loss

            # if loss > 2.5 * mbs:
            #     print("Bad initializations")
            #     sys.exit()
            
            loss.backward()

            model.optimizer.step()

            epoch_loss += loss.item()
            t2 = time.time()
            batch_time = t2-t1
            print ("batch = ", count, "loss = ", loss.item(), "time = ", np.round(batch_time, 2))
            model.optimizer.zero_grad() 
            
            loss = 0
        
        epoch_loss_list.append(epoch_loss)
        print('** Epoch loss is = ', epoch_loss)
        
        if epoch_loss < current_best_loss:
            torch.save(model.state_dict(), path+"/model_epoch_best.pth")
            torch.save(model.state_dict(), path+"/model_epoch"+str(epoch)+".pth")   
            current_best_loss = epoch_loss
            torch.save(epoch_loss_list, path+"/loss_list.pth")
        epoch_loss = 0

        
        
        print("-----------------")
        

    torch.save(train_args, path+"/training_config.pth")
    torch.save(epoch_loss_list, path+"/loss_list.pth")
    
    # plt.plot(epoch_loss_list)
    # plt.xlabel('Iteration')
    # plt.ylabel('Loss')
    # plt.yscale('log')
    # plt.title('Loss vs. Iteration')
    # plt.show()