<commit_before>import argparse, json, os, time

from parcellearning import gatpen
from parcellearning.utilities import gnnio
from parcellearning.utilities.early_stop import EarlyStopping
from parcellearning.utilities.batch import partition_graphs
from parcellearning.utilities.load import load_schema

from shutil import copyfile
from pathlib import Path

import pandas as pd
import numpy as np

import dgl
from dgl.data import register_data_args
import dgl.function as fn

import torch
import torch.nn.functional as F


def main(args):

    schema = load_schema(args.schema_file)
    in_dir = schema['data']['in']
    out_dir = schema['data']['out']
    Path(out_dir).mkdir(parents=True, exist_ok=True)

    # copy schema file to output directory
    copy_schema = ''.join([out_dir, args.schema_file.split('/')[-1]])
    if not os.path.exists(copy_schema):
        copyfile(args.schema_file, copy_schema)

    # get features
    features = schema['features']
    features.sort()

    ##### GET PARAMETERS FROM SCHEMA FILE #####
    # - - - - - - - - - - - - - - - - - - - - #
    # - - - - - - - - - - - - - - - - - - - - #

    MODEL_PARAMS = schema['model_parameters']
    OPT_PARAMS = schema['optimizer_parameters']
    TRAIN_PARAMS = schema['training_parameters']
    STOP_PARAMS = schema['stopping_parameters']

    REG = pd.read_csv(schema['spatial_regularizer'], index_col=[0])
    REG = torch.Tensor(np.asarray(REG))

    # - - - - - - - - - - - - - - - - - - - - #
    # - - - - - - - - - - - - - - - - - - - - #

    # load training and validation data
    training = gnnio.dataset(dType='training',
                             features=features,
                             dSet=schema['data']['training'],
                             norm=True,
                             aggregate=True,
                             clean=True)

    validation = gnnio.dataset(dType='validation', 
                               features=features,
                               dSet=schema['data']['validation'],
                               norm=True,
                               aggregate=True,
                               clean=True)

    validation = dgl.batch(validation)

    val_X = validation.ndata['features']
    val_Y = validation.ndata['label']

    ##### MODEL TRAINING #####
    # - - - - - - - - - - - - #
    # - - - - - - - - - - - - #

    # instantiate model using schema parameters
    model = gatpen.GATPEN(**MODEL_PARAMS)

    # instantiate Adam optimizer using scheme parameters
    optimizer = torch.optim.Adam(model.parameters(), **OPT_PARAMS)

    # initialize early stopper
    stopped_model_output='%s%s.earlystop.Loss.pt' % (out_dir, schema['model'])
    stopper = EarlyStopping(filename=stopped_model_output, **STOP_PARAMS)

    progress = {k: [] for k in ['Epoch',
                                'Duration',
                                'Train Loss',
                                'Train Acc',
                                'Val Loss',
                                'Val Acc']}

    cross_entropy = torch.nn.CrossEntropyLoss()

    dur = []
    print('\nTraining model\n')
    for epoch in range(TRAIN_PARAMS['epochs']):

        # learn model on training data
        batches = partition_graphs(training, TRAIN_PARAMS['n_batch'])

        model.train()
        t0 = time.time()

        # zero the gradients for this epoch
        optimizer.zero_grad()
        
        # aggregate training batch losses
        train_loss = 0

        # aggregate training batch accuracies
        train_acc = 0

        for iteration, batch in enumerate(batches):

            # get training features for this batch
            batch_X = batch.ndata['features']
            batch_Y = batch.ndata['label']

            # push batch through network
            batch_logits = model(batch, batch_X, **{'cost': REG})
            
            # compute batch performance
            # loss
            batch_loss = cross_entropy(batch_logits, batch_Y)
            
            # accuracy
            _, batch_indices = torch.max(F.softmax(batch_logits, dim=1), dim=1)
            batch_acc = (batch_indices == batch_Y).sum() / batch_Y.shape[0]

            # apply backward parameter update pass
            batch_loss.backward()

            print('Batch: %i | Batch Acc: %.3f | Batch Loss: %.3f ' % (iteration+1, batch_acc.item(), batch_loss.item()))

            # update training performance
            train_loss += batch_loss
            train_acc += batch_acc

            # accumulate the gradients from each batch
            if (iteration+1) % TRAIN_PARAMS['n_batch'] == 0:
                optimizer.step()
                optimizer.zero_grad()

        dur.append(time.time() - t0)

        # switch model into evaluation mode 
        # so we don't update the gradients using the validation data
        model.eval()
        with torch.no_grad():
            # push validation through network
            val_logits = model(validation, val_X)

            # compute validation performance
            # loss
            val_loss = cross_entropy(val_logits, val_Y, **{'cost': REG})

            # accuracy
            _, val_indices = torch.max(F.softmax(val_logits, dim=1), dim=1)
            val_acc = (val_indices == val_Y).sum() / val_Y.shape[0]

        train_loss /= TRAIN_PARAMS['n_batch']
        train_acc /= TRAIN_PARAMS['n_batch']

        # Show current performance
        print("Epoch {:05d} | Time(s) {:.4f} | Train Loss {:.4f} | Train Acc {:.4f} | Val Loss {:.4f} | Val Acc {:.4f}".format(
            epoch, np.mean(dur),
            train_loss.item(), train_acc.item(),
            val_loss.item(), val_acc.item()))
<commit_msg>GAT with spatial regularization penalty<commit_after>import argparse, json, os, time

from parcellearning import gatpen
from parcellearning.utilities import gnnio
from parcellearning.utilities.early_stop import EarlyStopping
from parcellearning.utilities.batch import partition_graphs
from parcellearning.utilities.load import load_schema

from shutil import copyfile
from pathlib import Path

import pandas as pd
import numpy as np

import dgl
from dgl.data import register_data_args
import dgl.function as fn

import torch
import torch.nn.functional as F


def main(args):

    schema = load_schema(args.schema_file)
    in_dir = schema['data']['in']
    out_dir = schema['data']['out']
    Path(out_dir).mkdir(parents=True, exist_ok=True)

    # copy schema file to output directory
    copy_schema = ''.join([out_dir, args.schema_file.split('/')[-1]])
    if not os.path.exists(copy_schema):
        copyfile(args.schema_file, copy_schema)

    # get features
    features = schema['features']
    features.sort()

    ##### GET PARAMETERS FROM SCHEMA FILE #####
    # - - - - - - - - - - - - - - - - - - - - #
    # - - - - - - - - - - - - - - - - - - - - #

    MODEL_PARAMS = schema['model_parameters']
    OPT_PARAMS = schema['optimizer_parameters']
    TRAIN_PARAMS = schema['training_parameters']
    STOP_PARAMS = schema['stopping_parameters']

    REG = pd.read_csv(schema['spatial_regularizer'], index_col=[0])
    REG = torch.Tensor(np.asarray(REG))

    # - - - - - - - - - - - - - - - - - - - - #
    # - - - - - - - - - - - - - - - - - - - - #

    # load training and validation data
    training = gnnio.dataset(dType='training',
                             features=features,
                             dSet=schema['data']['training'],
                             norm=True,
                             aggregate=True,
                             clean=True)

    validation = gnnio.dataset(dType='validation', 
                               features=features,
                               dSet=schema['data']['validation'],
                               norm=True,
                               aggregate=True,
                               clean=True)

    validation = dgl.batch(validation)

    val_X = validation.ndata['features']
    val_Y = validation.ndata['label']

    ##### MODEL TRAINING #####
    # - - - - - - - - - - - - #
    # - - - - - - - - - - - - #

    # instantiate model using schema parameters
    model = gatpen.GATPEN(**MODEL_PARAMS)

    # instantiate Adam optimizer using scheme parameters
    optimizer = torch.optim.Adam(model.parameters(), **OPT_PARAMS)

    # initialize early stopper
    stopped_model_output='%s%s.earlystop.Loss.pt' % (out_dir, schema['model'])
    stopper = EarlyStopping(filename=stopped_model_output, **STOP_PARAMS)

    progress = {k: [] for k in ['Epoch',
                                'Duration',
                                'Train Loss',
                                'Train Acc',
                                'Val Loss',
                                'Val Acc']}

    cross_entropy = torch.nn.CrossEntropyLoss()

    dur = []
    print('\nTraining model\n')
    for epoch in range(TRAIN_PARAMS['epochs']):

        # learn model on training data
        batches = partition_graphs(training, TRAIN_PARAMS['n_batch'])

        model.train()
        t0 = time.time()

        # zero the gradients for this epoch
        optimizer.zero_grad()
        
        # aggregate training batch losses
        train_loss = 0

        # aggregate training batch accuracies
        train_acc = 0

        for iteration, batch in enumerate(batches):

            # get training features for this batch
            batch_X = batch.ndata['features']
            batch_Y = batch.ndata['label']

            # push batch through network
            batch_logits = model(batch, batch_X, **{'cost': REG})
            # compute softmax of network
            batch_SM = F.softmax(batch_logits, dim=1)
            
            # compute batch performance
            # loss
            batch_loss = cross_entropy(batch_SM, batch_Y)
            
            # accuracy
            _, batch_indices = torch.max(batch_SM, dim=1)
            batch_acc = (batch_indices == batch_Y).sum() / batch_Y.shape[0]

            # apply backward parameter update pass
            batch_loss.backward()

            print('Batch: %i | Batch Acc: %.3f | Batch Loss: %.3f ' % (iteration+1, batch_acc.item(), batch_loss.item()))

            # update training performance
            train_loss += batch_loss
            train_acc += batch_acc

            # accumulate the gradients from each batch
            if (iteration+1) % TRAIN_PARAMS['n_batch'] == 0:
                optimizer.step()
                optimizer.zero_grad()

        dur.append(time.time() - t0)

        # switch model into evaluation mode 
        # so we don't update the gradients using the validation data
        model.eval()
        with torch.no_grad():
            # push validation through network
            val_logits = model(validation, val_X)
            val_SM = F.softmax(val_logits, dim=1)

            # compute validation performance
            # loss
            val_loss = cross_entropy(val_SM, val_Y, **{'cost': REG})

            # accuracy
            _, val_indices = torch.max(val_SM, dim=1)
            val_acc = (val_indices == val_Y).sum() / val_Y.shape[0]

        train_loss /= TRAIN_PARAMS['n_batch']
        train_acc /= TRAIN_PARAMS['n_batch']

        # Show current performance
        print("Epoch {:05d} | Time(s) {:.4f} | Train Loss {:.4f} | Train Acc {:.4f} | Val Loss {:.4f} | Val Acc {:.4f}".format(
            epoch, np.mean(dur),
            train_loss.item(), train_acc.item(),
            val_loss.item(), val_acc.item()))

{'fraction_non_alphanumeric': 0.1199026171637249, 'fraction_numerical': 0.004173550126075993, 'mean_word_length': 1.813600782778865, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7462321', 'n_tokens_mistral': 3427, 'n_tokens_neox': 3282, 'n_words': 732}