get_main.py

import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F
from sklearn.model_selection import train_test_split
import os
from termcolor import colored
from sklearn.metrics import brier_score_loss

import import_data as impt
from class_DeepHit import Model_DeepHit
from utils_eval import weighted_c_index, weighted_brier_score

_EPSILON = 1e-08


##### USER-DEFINED FUNCTIONS #####
def log(x):
    return torch.log(x + _EPSILON)


def div(x, y):
    return x / (y + _EPSILON)


def f_get_minibatch(mb_size, x, label, time, mask1, mask2):
    idx = np.random.choice(np.arange(np.shape(x)[0]), mb_size, replace=False)

    x_mb = x[idx, :].astype(np.float32)
    k_mb = label[idx, :].astype(np.float32)  # censoring(0)/event(1,2,..) label
    t_mb = time[idx, :].astype(np.float32)
    m1_mb = mask1[idx, :, :].astype(np.float32)  # fc_mask
    m2_mb = mask2[idx, :].astype(np.float32)  # fc_mask
    return (
        torch.tensor(x_mb),
        torch.tensor(k_mb),
        torch.tensor(t_mb),
        torch.tensor(m1_mb),
        torch.tensor(m2_mb),
    )


def get_valid_performance(
    DATA,
    MASK,
    in_parser,
    out_itr,
    eval_time=None,
    MAX_VALUE=-99,
    OUT_ITERATION=5,
    seed=1234,
):
    ##### DATA & MASK
    (data, time, label) = DATA
    (mask1, mask2) = MASK

    x_dim = np.shape(data)[1]
    _, num_Event, num_Category = np.shape(
        mask1
    )  # dim of mask1: [subj, Num_Event, Num_Category]

    ACTIVATION_FN = {"relu": F.relu, "elu": F.elu, "tanh": torch.tanh}

    ##### HYPER-PARAMETERS
    mb_size = in_parser["mb_size"]
    iteration = in_parser["iteration"]
    keep_prob = in_parser["keep_prob"]
    lr_train = in_parser["lr_train"]

    alpha = in_parser["alpha"]  # for log-likelihood loss
    beta = in_parser["beta"]  # for ranking loss
    gamma = in_parser["gamma"]  # for RNN-prediction loss

    parameter_name = (
        "a"
        + str("%02.0f" % (10 * alpha))
        + "b"
        + str("%02.0f" % (10 * beta))
        + "c"
        + str("%02.0f" % (10 * gamma))
    )

    # Xavier initializer is GlorotUniform in TensorFlow 2.x
    initial_W = torch.nn.init.xavier_uniform_

    ##### MAKE DICTIONARIES
    # INPUT DIMENSIONS
    input_dims = {"x_dim": x_dim, "num_Event": num_Event, "num_Category": num_Category}

    # NETWORK HYPER-PARAMETERS
    network_settings = {
        "h_dim_shared": in_parser["h_dim_shared"],
        "num_layers_shared": in_parser["num_layers_shared"],
        "h_dim_CS": in_parser["h_dim_CS"],
        "num_layers_CS": in_parser["num_layers_CS"],
        "active_fn": ACTIVATION_FN[in_parser["active_fn"]],
        "initial_W": initial_W,
        "keep_prob": keep_prob,
    }

    file_path_final = in_parser["out_path"] + "/itr_" + str(out_itr)

    # Create directories if they don't exist
    os.makedirs(file_path_final + "/models/", exist_ok=True)

    print(file_path_final + " (a:" + str(alpha) + " b:" + str(beta) + ")")

    ##### CREATE DEEPHIT NETWORK
    model = Model_DeepHit(input_dims, network_settings)
    optimizer = optim.Adam(model.parameters(), lr=lr_train)
    ### TRAINING-TESTING SPLIT
    (
        tr_data,
        te_data,
        tr_time,
        te_time,
        tr_label,
        te_label,
        tr_mask1,
        te_mask1,
        tr_mask2,
        te_mask2,
    ) = train_test_split(
        data, time, label, mask1, mask2, test_size=0.20, random_state=seed
    )

    (
        tr_data,
        va_data,
        tr_time,
        va_time,
        tr_label,
        va_label,
        tr_mask1,
        va_mask1,
        tr_mask2,
        va_mask2,
    ) = train_test_split(
        tr_data,
        tr_time,
        tr_label,
        tr_mask1,
        tr_mask2,
        test_size=0.20,
        random_state=seed,
    )

    # Fit normalization on the TRAINING split only (no leakage), then apply the
    # same parameters to validation (and, at eval time, to test).
    norm_params = impt.f_get_norm_params(tr_data, "standard")
    tr_data = impt.f_apply_Normalization(tr_data, norm_params)
    va_data = impt.f_apply_Normalization(va_data, norm_params)

    # Convert va_data to a tensor
    va_data = torch.tensor(
        va_data, dtype=torch.float32
    )  # ensure the data is a PyTorch tensor

    max_valid = -99
    stop_flag = 0

    if eval_time is None:
        eval_time = [
            int(np.percentile(tr_time, 25)),
            int(np.percentile(tr_time, 50)),
            int(np.percentile(tr_time, 75)),
        ]

    ### MAIN TRAINING LOOP
    print("MAIN TRAINING ...")
    print("EVALUATION TIMES: " + str(eval_time))

    avg_loss = 0
    for itr in range(iteration):
        if stop_flag > 5:  # Early stopping condition
            break

        # Fetch minibatch
        x_mb, k_mb, t_mb, m1_mb, m2_mb = f_get_minibatch(
            mb_size, tr_data, tr_label, tr_time, tr_mask1, tr_mask2
        )
        DATA = (x_mb, k_mb, t_mb)
        MASK = (m1_mb, m2_mb)
        PARAMETERS = (alpha, beta, gamma)

        # Train the model on the current batch
        loss_curr = model.training_step(DATA, MASK, PARAMETERS, optimizer)
        avg_loss += loss_curr / 1000

        if (itr + 1) % 1000 == 0:
            print(
                "|| ITR: "
                + str("%04d" % (itr + 1))
                + " | Loss: "
                + colored(str("%.4f" % (avg_loss)), "yellow", attrs=["bold"])
            )
            avg_loss = 0
            ### VALIDATION (based on average C-index of our interest) ###
            model.eval()  # Set the model to evaluation mode
            with torch.no_grad():
                pred = model(va_data)

            ### EVALUATION ###
            va_result1 = np.zeros([num_Event, len(eval_time)])

            for t, t_time in enumerate(eval_time):
                eval_horizon = int(t_time)

                if eval_horizon >= num_Category:
                    print("ERROR: evaluation horizon is out of range")
                    va_result1[:, t] = -1
                else:
                    risk = torch.sum(
                        pred[:, :, : (eval_horizon + 1)], dim=2
                    ).numpy()  # risk score until eval_time
                    for k in range(num_Event):
                        va_result1[k, t] = weighted_c_index(
                            tr_time,
                            (tr_label[:, 0] == k + 1).astype(int),
                            risk[:, k],
                            va_time,
                            (va_label[:, 0] == k + 1).astype(int),
                            eval_horizon,
                        )

            tmp_valid = np.mean(va_result1)

            if tmp_valid > max_valid:
                stop_flag = 0
                max_valid = tmp_valid
                print(f"Updated... Average C-index = {tmp_valid:.4f}")
                if max_valid > MAX_VALUE:
                    torch.save(
                        model.state_dict(),
                        os.path.join(
                            file_path_final, "models", f"model_itr_{out_itr}.pth"
                        ),
                    )

            else:
                stop_flag += 1

    return max_valid