py_mgcn/code/train.py

from __future__ import division
from __future__ import print_function

import time
import argparse

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

from sklearn.metrics import f1_score

from utils import load_data, class_accuracy, class_f1, layer_accuracy, layer_f1, dict_to_writer, trains_vals_tests_split
from models import CCN
from tensorboardX import SummaryWriter

# f1 name should be f1_micro or f1_macro
# class_metrics = {'loss': F.nll_loss, 'acc': class_accuracy, 'f1_micro': class_f1, 'f1_macro': class_f1,
#                  'all_f1_micro': f1_score, 'all_f1_macro': f1_score}
class_metrics = {'loss': F.nll_loss, 'all_f1_micro': f1_score, 'all_f1_macro': f1_score}
layer_metrics = {'loss': torch.nn.BCEWithLogitsLoss()}


# Preparing the output of classifier for every layer separately (class numbers should be different over layers)
def prepare_output_labels(outputs, labels, idx):
    tmp_labels = []
    tmp_outputs = []
    max_label = 0
    for i in range(len(labels)):
        tmp_labels.extend(labels[i][idx[i]] + max_label)
        tmp_outputs.extend((outputs[i].max(1)[1].type_as(labels[i]) + max_label)[idx[i]])
        max_label = labels[i].max()

    return tmp_outputs, tmp_labels


# Calculate classification metrics for every layer
def model_stat_in_layer_class(in_classes_pred, in_labels, idx, metrics):
    stats = {}
    for i in range(len(in_classes_pred)):
        if not in_classes_pred[i] is None:
            for metr_name,metr_func in metrics.items():
                if not metr_name in stats:
                    stats[metr_name] = []
                if metr_name[:2] == 'f1':
                    try:
                        stats[metr_name].append(metr_func(in_classes_pred[i][idx[i]], in_labels[i][idx[i]], metr_name.split('f1_')[-1]))
                    except:
                        print("Exception in F1 for class" + str(i) + " for classification")
                        stats[metr_name].append(0)
                elif metr_name[:3] == 'all':
                    continue
                else:
                    stats[metr_name].append(metr_func(in_classes_pred[i][idx[i]], in_labels[i][idx[i]]))
        else:
            for metr_name, metr_func in metrics.items():
                if not metr_name in stats:
                    stats[metr_name] = []
                stats[metr_name].append(0)

    tmp_in_classes_pred, tmp_labels = prepare_output_labels(in_classes_pred, in_labels, idx)
    for metr_name, metr_func in metrics.items():
        if metr_name[:3] == 'all':
            stats[metr_name] = metr_func(tmp_in_classes_pred, tmp_labels, average=metr_name.split('f1_')[-1])

    return stats


# Calculate reconstruction metrics for within and between layer edges (based on the input)
def model_stat_struc(preds, reals, metrics, pos_weights, norms, idx=None):
    stats = {}
    for i in range(len(preds)):
        if reals[i] is None:
            continue
        if 'sparse' in reals[i].type():
            curr_real = reals[i].to_dense()
        else:
            curr_real = reals[i]
        if idx is None:
            final_preds = preds[i]
            final_real = curr_real
        else:
            final_preds = preds[i][idx[i][0],idx[i][1]]
            final_real = curr_real[idx[i][0],idx[i][1]]
        for metr_name,metr_func in metrics.items():
            if not metr_name in stats:
                stats[metr_name] = []
            if metr_name == 'f1':
                try:
                    stats[metr_name].append(metr_func(final_preds, final_real, type = metr_name.split('f1_')[-1]))
                except:
                    print("Exception in F1 for layer" + str(i) + " for structure")
                    stats[metr_name].append(0)
            elif metr_name == 'loss':
                if args.cuda:
                    pw = pos_weights[i] * torch.ones(final_real.size()).cuda()
                else:
                    pw = pos_weights[i] * torch.ones(final_real.size())
                stats[metr_name].append(norms[i]*torch.nn.BCEWithLogitsLoss(pos_weight=pw)(final_preds, final_real))
            else:
                try:
                    stats[metr_name].append(metr_func(final_preds, final_real))
                except:
                    print("Odd problem")

    return stats


def train(epoch):
    t = time.time()
    model.train()
    optimizer.zero_grad()

    # Run model
    t1 = time.time()
    classes_pred, bet_layers_pred, in_layers_pred = model(features, adjs)
    # Final statistics
    stats = dict()
    stats['in_class'] = model_stat_in_layer_class(classes_pred, labels, idx_trains, class_metrics)
    # Reconstructing all the within and between layer edges without train and test split (for using part of the edges
    # for reconstruction, idx is needed)
    stats['in_struc'] = model_stat_struc(in_layers_pred, adjs_orig, layer_metrics, adjs_pos_weights, adjs_norms)
    stats['bet_struc'] = model_stat_struc(bet_layers_pred, bet_adjs_orig, layer_metrics, bet_pos_weights, bet_norms)
    # Write to writer
    for name, stat in stats.items():
        dict_to_writer(stat, writer, epoch, name, 'Train')
    # Calculate loss
    loss_train = 0
    for name, stat in stats.items():
        # Weighted loss
        if name == 'in_struc':
            loss_train += wlambda * sum([a * b for a, b in zip(stat['loss'], adj_weight)])
        elif name== 'bet_struc':
            loss_train += wlambda * sum([a * b for a, b in zip(stat['loss'], adj_weight)])
        else:
            loss_train += sum([a * b for a, b in zip(stat['loss'], adj_weight)])

    # Gradients
    loss_train.backward()
    optimizer.step()

    # Validation if needed
    if not args.fastmode:
        model.eval()
        classes_pred, bet_layers_pred, in_layers_pred = model(features, adjs)
        stats = dict()
        stats['in_class'] = model_stat_in_layer_class(classes_pred, labels, idx_vals, class_metrics)

        stats['in_struc'] = model_stat_struc(in_layers_pred, adjs_orig, layer_metrics, adjs_pos_weights, adjs_norms)
        stats['bet_struc'] = model_stat_struc(bet_layers_pred, bet_adjs_orig, layer_metrics, bet_pos_weights, bet_norms)

        for name, stat in stats.items():
            dict_to_writer(stat, writer, epoch, name, 'Validation')
        loss_val = sum([sum(stat['loss']) for stat in stats.values()])

    print('Epoch: {:04d}'.format(epoch+1),
          'loss_train: {:.4f}'.format(loss_train.item()),
          'time: {:.4f}s'.format(time.time() - t))


def test():
    # Test
    model.eval()
    classes_pred, bet_layers_pred, in_layers_pred = model(features, adjs)
    stats = dict()
    stats['in_class'] = model_stat_in_layer_class(classes_pred, labels, idx_tests, class_metrics)

    for name, stat in stats.items():
        dict_to_writer(stat, writer, epoch, name, 'Test')

    loss_test = sum([sum(stat['loss']) for stat in stats.values()])
    print("Test set results:",
          "loss= {:.4f}".format(loss_test.item()))


if __name__=='__main__':
    # Training settings
    parser = argparse.ArgumentParser()
    parser.add_argument('--no-cuda', action='store_true', default=True,
                        help='Disables CUDA training.')
    parser.add_argument('--fastmode', action='store_true', default=False,
                        help='Validate during training pass.')
    parser.add_argument('--seed', type=int, default=42, help='Random seed.')
    parser.add_argument('--epochs', type=int, default=20,
                        help='Number of epochs to train.')
    parser.add_argument('--weight_decay', type=float, default=5e-4,
                        help='Weight decay (L2 loss on parameters).')
    parser.add_argument('--dropout', type=float, default=0.5,
                        help='Dropout rate (1 - keep probability).')

    args = parser.parse_args()
    args.cuda = not args.no_cuda and torch.cuda.is_available()

    dataset_str = "infra"

    # parameter
    # All combination of parameters will be tested
    # adj_weights contains a list of different configs, for example [[1,1,1]] contains one config with equal weights
    # for a three-layered input
    adj_weights = []
    # parameter, weight of link reconstruction loss function
    wlambdas = [10]
    # parameter, hidden dimension for every layer should be defined
    hidden_structures = [[[32],[32],[32]]]
    # hidden_structures = [[[16], [16], [16]],[[32], [32], [32]],[[64], [64], [64]],[[128], [128], [128]],[[256], [256], [256]]]
    # Learning rate
    lrs = [0.01]
    # test size
    test_sizes = [0.2]

    # Load data
    adjs, adjs_orig, adjs_sizes, adjs_pos_weights, adjs_norms, bet_pos_weights, bet_norms, bet_adjs, bet_adjs_orig, bet_adjs_sizes, \
    features, features_sizes, labels, labels_nclass = load_data(path="../data/" + dataset_str + "/", dataset=dataset_str)
    # Number of layers
    n_inputs = len(adjs)
    # Weights of layers
    adj_weights.append(n_inputs * [1])

    if args.cuda:
        for i in range(n_inputs):
            labels[i] = labels[i].cuda()

            adjs_pos_weights[i] = adjs_pos_weights[i].cuda()
            adjs[i] = adjs[i].cuda()
            adjs_orig[i] = adjs_orig[i].cuda()
            features[i] = features[i].cuda()

        for i in range(len(bet_adjs)):
            if not bet_adjs[i] is None:
                bet_adjs[i] = bet_adjs[i].cuda()
                bet_adjs_orig[i] = bet_adjs_orig[i].cuda()
                bet_pos_weights[i] = bet_pos_weights[i].cuda()
    # Number of runs
    for run in range(10):
        for ts in test_sizes:
            idx_trains, idx_vals, idx_tests = trains_vals_tests_split(n_inputs, [s[0] for s in adjs_sizes], val_size=0.1,
                                                                      test_size=ts, random_state=int(run + ts*100))
            if args.cuda:
                for i in range(n_inputs):
                    idx_trains[i] = idx_trains[i].cuda()
                    idx_vals[i] = idx_vals[i].cuda()
                    idx_tests[i] = idx_tests[i].cuda()
            # Train model
            t_total = time.time()
            for wlambda in wlambdas:
                for adj_weight in adj_weights:
                    for lr in lrs:
                        for hidden_structure in hidden_structures:
                            temp_weight = ['{:.2f}'.format(x) for x in adj_weight]
                            w_str = '-'.join(temp_weight)
                            h_str = '-'.join([','.join(map(str, temp)) for temp in hidden_structure])
                            writer = SummaryWriter('log/' + dataset_str + '_run-'+ str(run)+ '/' + '_lambda-' +
                                                   str(wlambda) + "_adj_w-" + w_str + "_LR-" + str(lr) +
                                                             '_hStruc-' + h_str + '_test_size-' + str(ts))
                            model = CCN(n_inputs=n_inputs,
                                        inputs_nfeat=features_sizes,
                                        inputs_nhid=hidden_structure,
                                        inputs_nclass=labels_nclass,
                                        dropout=args.dropout)
                            optimizer = optim.Adam(model.parameters(),
                                                   lr=lr, weight_decay=args.weight_decay)
                            if args.cuda:
                                model.cuda()
                            for epoch in range(args.epochs):
                                train(epoch)
                                # Testing
                                test()
                            print("Optimization Finished!")
                            print("Total time elapsed: {:.4f}s".format(time.time() - t_total))