MONBM_Problem.py

# -*- coding: utf-8 -*-
import numpy as np
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

from itertools import product
from sklearn.metrics import roc_auc_score

import geatpy as ea
import time

def cal_DP(y_pred, s_batch, requires_grad=True):
    if requires_grad:
        y_pred.requires_grad = True
        s_batch.requires_grad = True

    # 计算Demographic Parity
    indices1 = torch.nonzero(s_batch == 1).squeeze()
    indices2 = torch.nonzero(s_batch == 0).squeeze()
    temp1 = torch.mean(y_pred[indices1])
    temp2 = torch.mean(y_pred[indices2])
    dp = torch.abs(temp1 - temp2)
    return dp

def cal_complexity_uc(y_out_feats):
    # if requires_grad:
    #     x.requires_grad = True
    # y_out, attribution = model.forward_cal(x)
    fraction = torch.abs(y_out_feats) / torch.sum(torch.abs(y_out_feats), dim=1).unsqueeze(1)
    temp = fraction * torch.log(fraction)
    complexity_uc = torch.mean(-torch.nansum(temp, dim=1))
    return complexity_uc

def cal_GU(model, x, flag, k, k1, k2):
    points = np.linspace(0, 1, k)
    loss = torch.tensor(0.0, requires_grad=True)
    pred_data = torch.tensor(np.repeat(points, x.shape[1]).reshape(-1, x.shape[1])).float()
    feature_contribution = model.feature_contribution(pred_data)
    max_feature_contribution = torch.max(feature_contribution, dim=0)[0]
    min_feature_contribution = torch.min(feature_contribution, dim=0)[0]

    for idx in range(x.shape[1]):
        if flag[idx] == 0:
            dy = torch.abs(feature_contribution[1:, idx] - feature_contribution[:-1, idx]) / (max_feature_contribution[idx] - min_feature_contribution[idx]) - k2[idx]
            temp_loss = torch.mean(torch.relu(dy))
            loss = loss + temp_loss
        elif flag[idx] == 1:
            dy = (feature_contribution[1:, idx] - feature_contribution[:-1, idx]) / (max_feature_contribution[idx] - min_feature_contribution[idx])
            temp_loss1 = torch.mean(torch.relu(k1[idx] - dy))
            temp_loss2 = torch.mean(torch.relu(dy - k2[idx]))
            loss = loss + temp_loss1 + temp_loss2
        elif flag[idx] == 2:
            dy = (feature_contribution[:-1, idx] - feature_contribution[1:, idx]) / (max_feature_contribution[idx] - min_feature_contribution[idx])
            temp_loss1 = torch.mean(torch.relu(k1[idx] - dy))
            temp_loss2 = torch.mean(torch.relu(dy - k2[idx]))
            loss = loss + temp_loss1 + temp_loss2
    return loss

def cal_complexity_GI(y_out_feats):
    abs_y_out_feats = torch.abs(y_out_feats)
    sorted_abs_y_out_feats, sorted_idx = torch.sort(abs_y_out_feats, dim=1)
    sum_abs_y_out_feats = torch.sum(abs_y_out_feats, dim=1)
    frac_abs_y_out_feats = sorted_abs_y_out_feats / sum_abs_y_out_feats.unsqueeze(1)
    for k in range(y_out_feats.shape[1]):
        frac_abs_y_out_feats[:, k] = frac_abs_y_out_feats[:, k] * ((y_out_feats.shape[1] - k - 0.5) / y_out_feats.shape[1])
    result = torch.nanmean(1 - 2 * torch.sum(frac_abs_y_out_feats, dim=1))
    return result

def cal_LU(y_out_feats):
    uC = cal_complexity_uc(y_out_feats)
    GI = 1 - cal_complexity_GI(y_out_feats)
    g_mean = torch.pow(uC * GI, 1 / 2)
    return g_mean


class MONBMProblem(ea.Problem):
    def __init__(self, M=3, learning_rate=0.01, batch_size=500, sensitive_attribute=None, list_objs=None,
                 weight_decay=1e-3, Dim=None, data_name=None, feature_names=None, categorical_columns=None,
                 nbm_model=None, epochs=1, train_dataset=None, val_dataset=None, test_dataset=None, loss_fn=None,
                 monotonic_flag=None, monotonic_k=None, all_optimized=False, k1=None, k2=None):
        problem_name = 'MONBM_Problem'
        maxormins = [1] * M # 目标最大最小化
        varTypes = [] # 决策变量连续还是离散
        for feature_name in feature_names:
            if feature_name in categorical_columns:
                varTypes.append(1)
            else:
                varTypes.append(0)

        lb = [0] * Dim
        ub = [1] * Dim
        lbin = [1] * Dim
        ubin = [1] * Dim

        ea.Problem.__init__(self, problem_name, M, maxormins, Dim, varTypes, lb, ub, lbin, ubin)

        self.M = M
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.sensitive_attribute = sensitive_attribute
        self.list_objs = list_objs
        self.weight_decay = weight_decay
        self.Dim = Dim
        self.feature_names = feature_names
        self.categorical_columns = categorical_columns
        self.dataname=data_name
        self.nbm_model = nbm_model
        self.epochs = epochs
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.test_dataset = test_dataset
        self.loss_fn = loss_fn

        train = TensorDataset(train_dataset.features, train_dataset.labels, train_dataset.s,
                              torch.arange(start=0, end=train_dataset.features.shape[0], step=1))
        val = TensorDataset(val_dataset.features, val_dataset.labels, val_dataset.s,
                              torch.arange(start=0, end=val_dataset.features.shape[0], step=1))
        test = TensorDataset(test_dataset.features, test_dataset.labels, test_dataset.s,
                              torch.arange(start=0, end=test_dataset.features.shape[0], step=1))
        self.train_loader = DataLoader(train, batch_size=self.batch_size, shuffle=True)
        self.val_loader = DataLoader(val, batch_size=self.batch_size, shuffle=True)
        self.test_loader = DataLoader(test, batch_size=self.batch_size, shuffle=True)
        self.monotonic_flag = monotonic_flag
        self.monotonic_k = monotonic_k
        self.all_optimized = all_optimized
        self.k1 = k1
        self.k2 = k2

    def getFeature(self):
        return self.num_features

    def partial_training(self, pop, loss_type=-1):
        pop_size = len(pop)
        # 优化器
        learning_rate = self.learning_rate
        weight_decay = self.weight_decay
        if self.loss_fn == 'BCEWithLogitsLoss':
            loss_fn = torch.nn.BCEWithLogitsLoss()

        for pop_i in np.arange(pop_size):
            individual = pop.Chrom[pop_i]
            individual_model = individual.model

            if loss_type != -1:
                individual.train()
                for i, (x_batch, y_batch, s_batch, data_idx) in enumerate(self.train_loader):
                    y_out, y_out_feats = individual_model(x_batch)
                    y_logits = torch.sigmoid_(y_out.clone())
                    y_pred = y_logits > 0.5
                    y_pred = y_pred.int()
                    if loss_type[pop_i] == 'BCE':
                        loss = loss_fn(y_out, y_batch.float())
                    elif loss_type[pop_i] == 'DP':
                        loss = cal_DP(y_pred.float(), s_batch.float())
                    elif loss_type[pop_i] == 'GU':
                        loss = cal_GU(individual_model, x_batch.float(), self.monotonic_flag,
                                                        self.monotonic_k, self.k1, self.k2)
                    elif loss_type[pop_i] == 'LU':
                        loss = cal_LU(y_out_feats)
                    params_to_update = []
                    params_to_update_count = 0
                    if self.all_optimized == False:
                        for name, param in individual_model.named_parameters():
                            if "featurizer" in name or "_bias" in name:
                                params_to_update.append(param)
                                params_to_update_count = params_to_update_count + param.numel()
                            else:
                                param.requires_grad = False
                    else:
                        for name, param in individual_model.named_parameters():
                            params_to_update.append(param)
                            params_to_update_count = params_to_update_count + param.numel()

                    optimizer = torch.optim.AdamW(
                        params_to_update,
                        lr=learning_rate,
                        weight_decay=weight_decay,
                    )
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()

            else:
                for epoch in np.arange(self.epochs):
                    individual_model.train()
                    for i, (x_batch, y_batch, s_batch, data_idx) in enumerate(self.train_loader):
                        y_out, y_out_feats = individual_model(x_batch)
                        y_logits = torch.sigmoid_(y_out.clone())
                        y_pred = y_logits > 0.5
                        y_pred = y_pred.int()

                        loss_list = []

                        if 'BCE' in self.list_objs or 'AUROC' in self.list_objs:
                            loss_acc = loss_fn(y_out, y_batch.float())
                            loss_list.append(loss_acc)
                        if 'DP' in self.list_objs:
                            loss_DP_classification = cal_DP(y_pred.float(), s_batch.float())
                            loss_list.append(loss_DP_classification)
                        if 'LU' in self.list_objs:
                            loss_com = cal_LU(y_out_feats)
                            loss_list.append(loss_com)
                        if 'GU' in self.list_objs:
                            loss_smooth = cal_GU(individual_model, x_batch.float(),
                                                      self.monotonic_flag, self.monotonic_k, self.k1, self.k2)
                            loss_list.append(loss_smooth)

                        loss = loss_list[np.random.permutation(len(loss_list))[0]]

                        params_to_update = []
                        params_to_update_count = 0
                        if self.all_optimized == False:
                            for name, param in individual_model.named_parameters():
                                if "featurizer" in name or "_bias" in name:
                                    params_to_update.append(param)
                                    params_to_update_count = params_to_update_count + param.numel()
                                else:
                                    param.requires_grad = False
                        else:
                            for name, param in individual_model.named_parameters():
                                params_to_update.append(param)
                                params_to_update_count = params_to_update_count + param.numel()

                        optimizer = torch.optim.AdamW(
                            params_to_update,
                            lr=learning_rate,
                            weight_decay=weight_decay,
                        )
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()

    def aimFunc(self, pop):  # 目标函数
        start = time.time()
        pop_size = len(pop)
        AllObj_train = np.zeros([pop_size, len(self.list_objs)])
        AllObj_val = np.zeros([pop_size, len(self.list_objs)])
        AllObj_test = np.zeros([pop_size, len(self.list_objs)])

        pred_label_train = np.zeros([pop_size, self.train_dataset.labels.shape[0]])
        pred_label_val = np.zeros([pop_size, self.val_dataset.labels.shape[0]])
        pred_label_test = np.zeros([pop_size, self.test_dataset.labels.shape[0]])

        pred_logits_train = np.zeros([pop_size, self.train_dataset.labels.shape[0]])
        pred_logits_val = np.zeros([pop_size, self.val_dataset.labels.shape[0]])
        pred_logits_test = np.zeros([pop_size, self.test_dataset.labels.shape[0]])

        for pop_i in np.arange(pop_size):
            individual = pop.Chrom[pop_i]
            individual_model = individual.model

            with torch.enable_grad():
                # individual.eval()

                # evaluate training dataset
                y_out, y_out_feats = individual_model.forward_cal(self.train_dataset.features)
                y_logits = torch.sigmoid_(y_out.clone())
                y_pred = y_logits > 0.5
                y_pred = y_pred.int()
                pred_logits_train[pop_i, :] = y_logits.detach().numpy().copy()
                pred_label_train[pop_i, :] = y_pred.detach().numpy().copy()
                for obj_i in np.arange(len(self.list_objs)):
                    if self.list_objs[obj_i] == 'BCE':
                        loss_fn = torch.nn.BCEWithLogitsLoss()
                        bce = loss_fn(y_out, self.train_dataset.labels.float())
                        AllObj_train[pop_i, obj_i] = bce
                    elif self.list_objs[obj_i] == 'DP':
                        DP = cal_DP(y_pred.float(), self.train_dataset.s,
                                                   requires_grad=False).numpy()
                        AllObj_train[pop_i, obj_i] = DP
                    elif self.list_objs[obj_i] == 'GU':
                        smooth = cal_GU(individual_model, self.train_dataset.features,
                                                  self.monotonic_flag, self.monotonic_k, self.k1, self.k2)
                        AllObj_train[pop_i, obj_i] = smooth
                    elif self.list_objs[obj_i] == 'LU':
                        LU = cal_LU(y_out_feats)
                        AllObj_train[pop_i, obj_i] = LU

                y_out, y_out_feats = individual_model.forward_cal(self.test_dataset.features)
                y_logits = torch.sigmoid_(y_out.clone())
                y_pred = y_logits > 0.5
                y_pred = y_pred.int()
                pred_logits_test[pop_i, :] = y_logits.detach().numpy().copy()
                pred_label_test[pop_i, :] = y_pred.detach().numpy().copy()
                for obj_i in np.arange(len(self.list_objs)):
                    if self.list_objs[obj_i] == 'BCE':
                        loss_fn = torch.nn.BCEWithLogitsLoss()
                        bce = loss_fn(y_out, self.test_dataset.labels.float())
                        AllObj_test[pop_i, obj_i] = bce
                    elif self.list_objs[obj_i] == 'DP':
                        DP = cal_DP(y_pred.float(), self.test_dataset.s,
                                                   requires_grad=False).numpy()
                        AllObj_test[pop_i, obj_i] = DP
                    elif self.list_objs[obj_i] == 'GU':
                        smooth = cal_GU(individual_model, self.test_dataset.features,
                                                  self.monotonic_flag, self.monotonic_k, self.k1, self.k2)
                        AllObj_test[pop_i, obj_i] = smooth
                    elif self.list_objs[obj_i] == 'LU':
                        LU = cal_LU(y_out_feats)
                        AllObj_test[pop_i, obj_i] = LU
        pop.CV = np.zeros([pop_size, 1])
        pop.ObjV = AllObj_train
        pop.ObjV_train = AllObj_train
        pop.ObjV_valid = AllObj_train
        pop.ObjV_test = AllObj_test

        pop.pred_label_train = pred_label_train
        pop.pred_label_valid = pred_label_train
        pop.pred_label_test = pred_label_test

        pop.pred_logits_train = pred_logits_train
        pop.pred_logits_valid = pred_logits_train
        pop.pred_logits_test = pred_logits_test

        return AllObj_train, AllObj_train