graphcl_model.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import global_add_pool
from torch_geometric.nn import GINEConv, GPSConv, TransformerConv, GINConv
import os
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score
from sklearn.manifold import TSNE
from collections import Counter
import evaluate_embedding as eval_emb

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class NeuralClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim=256):
        super(NeuralClassifier, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.LeakyReLU(0.1),  # 使用LeakyReLU代替ReLU
            nn.Dropout(0.3),    # 降低Dropout率
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.LeakyReLU(0.1),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        return self.classifier(x)

def train_neural_classifier(train_emb, train_labels, test_emb, test_labels, input_dim, device, epochs, batch_size):
    """
    训练神经网络分类器
    
    Args:
        train_emb: 训练集嵌入向量
        train_labels: 训练集标签
        test_emb: 测试集嵌入向量
        test_labels: 测试集标签
        input_dim: 输入维度
        device: 设备
        epochs: 训练轮数
        batch_size: 批次大小
    
    Returns:
        tuple: (测试集准确率, 测试集预测结果)
    """
    train_emb = torch.FloatTensor(train_emb).to(device)
    train_labels = torch.FloatTensor(train_labels).to(device)
    test_emb = torch.FloatTensor(test_emb).to(device)
    test_labels = torch.FloatTensor(test_labels).to(device)
    
    train_dataset = torch.utils.data.TensorDataset(train_emb, train_labels)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
    
    classifier = NeuralClassifier(input_dim).to(device)
    optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001)
    criterion = nn.BCEWithLogitsLoss()
    
    best_acc = 0
    best_preds = None
    for epoch in range(epochs):
        classifier.train()
        total_loss = 0
        for batch_emb, batch_labels in train_loader:
            optimizer.zero_grad()
            outputs = classifier(batch_emb)
            loss = criterion(outputs.squeeze(), batch_labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        # 评估
        if (epoch + 1) % 5 == 0:
            classifier.eval()
            with torch.no_grad():
                test_outputs = classifier(test_emb)
                test_preds = (test_outputs.squeeze() > 0.5).float()
                acc = (test_preds == test_labels).float().mean().item()
                if acc > best_acc:
                    best_acc = acc
                    best_preds = test_preds.cpu().numpy()
                print(f'Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}, Test Acc: {acc:.4f}')
    
    return best_acc, best_preds

def evaluate_model(model, dataloader_train, dataloader_test, output_dir, service, phase, epoch=None):
    """
    评估模型性能
    
    Args:
        model: 待评估的模型
        dataloader_train: 训练数据加载器
        dataloader_test: 测试数据加载器
        output_dir: 输出目录
        service: 服务名称
        phase: 评估阶段 (init/unsup/semi)
        epoch: 当前训练轮数（可选）
    
    Returns:
        dict: 包含评估结果的字典，以及t-SNE可视化所需的数据
    """
    model.eval()
    
    emb_train, y_train = model.encoder.get_embeddings(dataloader_train)
    emb_test, y_test = model.encoder.get_embeddings(dataloader_test)

    y_train_01 = y_train.copy()
    y_test_01 = y_test.copy()
    y_train_01[y_train_01 > 0] = 1
    y_test_01[y_test_01 > 0] = 1
    
    # 创建日志内容
    epoch_str = f"Epoch {epoch} - " if epoch is not None else ""
    
    # 保存评估结果
    with open(os.path.join(output_dir, f"evaluation_results_{service}.txt"), "a") as f:
        f.write(f"\n{epoch_str}{phase}\n")

    # 多分类器评估
    classification_results = {}
    
    # 添加神经网络分类器评估
    input_dim = emb_train.shape[1]
    nn_acc, nn_preds = train_neural_classifier(emb_train, y_train_01, emb_test, y_test_01, input_dim, device, epochs=50, batch_size=8)
    
    # 计算神经网络分类器的评估指标
    nn_precision = precision_score(y_test_01, nn_preds)
    nn_recall = recall_score(y_test_01, nn_preds)
    nn_f1 = f1_score(y_test_01, nn_preds)
    
    print(f'{epoch_str}{phase} - Model: Neural Network')
    print(f'  Test Acc: {nn_acc}')
    print(f'  Precision: {nn_precision}')
    print(f'  Recall: {nn_recall}')
    print(f'  F1 Score: {nn_f1}')
    
    classification_results['neural_network'] = {
        'test_acc': nn_acc,
        'precision': nn_precision,
        'recall': nn_recall,
        'f1': nn_f1
    }
    
    # 保存到评估结果文件
    with open(os.path.join(output_dir, f"evaluation_results_{service}.txt"), "a") as f:
        f.write(f'Model: Neural Network\n')
        f.write(f'  Test Acc: {nn_acc}\n')
        f.write(f'  Precision: {nn_precision}\n')
        f.write(f'  Recall: {nn_recall}\n')
        f.write(f'  F1 Score: {nn_f1}\n')
    
    # 其他分类器评估
    for eval_mode in ['svc', 'randomforest', 'linearsvc', 'logistic']:
        # 使用训练集训练分类器，在测试集上评估
        acc, preds = eval_emb.evaluate_embedding_with_split(emb_train, y_train_01, emb_test, y_test_01, eval_mode)
        
        # 计算评估指标
        precision = precision_score(y_test_01, preds)
        recall = recall_score(y_test_01, preds)
        f1 = f1_score(y_test_01, preds)
        
        print(f'{epoch_str}{phase} - Model: {eval_mode}')
        print(f'  Test Acc: {acc}')
        print(f'  Precision: {precision}')
        print(f'  Recall: {recall}')
        print(f'  F1 Score: {f1}')
        
        classification_results[eval_mode] = {
            'test_acc': acc,
            'precision': precision,
            'recall': recall,
            'f1': f1
        }
        
        # 保存到评估结果文件
        with open(os.path.join(output_dir, f"evaluation_results_{service}.txt"), "a") as f:
            f.write(f'Model: {eval_mode}\n')
            f.write(f'  Test Acc: {acc}\n')
            f.write(f'  Precision: {precision}\n')
            f.write(f'  Recall: {recall}\n')
            f.write(f'  F1 Score: {f1}\n')
    
    # t-SNE降维
    n_samples = emb_test.shape[0]
    perplexity = min(30, n_samples - 1) 
    
    tsne = TSNE(n_components=2, perplexity=perplexity)
    emb_2d = tsne.fit_transform(emb_test)
    
    # 返回结果字典，包含t-SNE数据
    return {
        'classification': classification_results,
        'tsne_data': {
            'emb_2d': emb_2d,
            'y_test': y_test,
            'phase': phase,
            'service': service
        }
    }

class Encoder(nn.Module):
    def __init__(self, num_features, num_edge_features, hidden_dim, num_gc_layers, global_dim, multi_heads_num=4):
        super(Encoder, self).__init__()

        self.global_dim = global_dim

        self.num_gc_layers = num_gc_layers
        self.initial_dim = 1536
        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()

        self.initial_projection = nn.Linear(num_features, self.initial_dim)

        for i in range(num_gc_layers):
            if i == 0:
                layer = nn.Sequential(nn.Linear(self.initial_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))
            else:
                layer = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim))
                

            conv = GINEConv(layer, edge_dim=num_edge_features)   # GINEConv可以考虑边特征
            bn = nn.BatchNorm1d(hidden_dim)
            self.convs.append(conv)
            self.bns.append(bn)
        
        # 针对图的全局时间特征做映射
        if global_dim > 0:
            self.global_projection = nn.Sequential(
                nn.Linear(1, global_dim),
                nn.ReLU()
            )

    def forward(self, x, edge_index, edge_attr, batch, time, debug=False):
        if x is None:
            x = torch.ones((batch.shape[0], 1)).to(device)
        if edge_attr is None:
            edge_attr = torch.zeros((edge_index.shape[1], 1)).to(device)

        x = self.initial_projection(x)

        xs = []
        for i in range(self.num_gc_layers):
            x = self.convs[i](x, edge_index, edge_attr=edge_attr)  # GINEConv
            x = F.relu(x)
            x = self.bns[i](x)
            xs.append(x)

        xpool = [global_add_pool(x, batch) for x in xs]
        x = torch.cat(xpool, 1)

        if self.global_dim > 0:
            x_time = self.global_projection(time.view(-1, 1))
            x = torch.cat([x, x_time], dim=1)

        return x, torch.cat(xs, 1)


    def get_embeddings(self, loader, debug=False):

        ret = []
        y = []
        with torch.no_grad():
            for data in loader:

                data = data[0]
                data.to(device)
                x, edge_index, edge_attr, batch = data.x, data.edge_index, data.edge_attr, data.batch
                if x is None:
                    x = torch.ones((batch.shape[0],1)).to(device)
                x, _ = self.forward(x, edge_index, edge_attr, batch, data.time, debug=debug)

                ret.append(x.cpu().numpy())
                y.append(data.y.cpu().numpy())
        ret = np.concatenate(ret, 0)
        y = np.concatenate(y, 0)
        return ret, y

class GraphCLModel(nn.Module):
    def __init__(self, hidden_dim, dataset_num_features, num_edge_features, num_gc_layers, global_dim,
                 alpha=0.5, beta=1., gamma=.1, multi_heads_num=4):
        super(GraphCLModel, self).__init__()

        assert hidden_dim % multi_heads_num == 0, "hidden_dim must be divisible by multi_heads_num"

        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.prior = False # default false

        self.embedding_dim = hidden_dim * num_gc_layers
        self.global_dim = global_dim
        self.encoder = Encoder(dataset_num_features, num_edge_features, hidden_dim, num_gc_layers, global_dim, multi_heads_num)

        self.proj_head = nn.Sequential(nn.Linear(self.embedding_dim + self.global_dim, self.embedding_dim), nn.ReLU(inplace=True), nn.Linear(self.embedding_dim, self.embedding_dim))

        self.init_emb()

    def init_emb(self):
        initrange = -1.5 / self.embedding_dim
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.fill_(0.0)

    def forward(self, data):
        x, edge_index, edge_attr, batch, num_graphs = data.x, data.edge_index, data.edge_attr, data.batch, data.num_graphs

        if x is None:
            x = torch.ones(batch.shape[0]).to(device)
        if edge_attr is None:
            edge_attr = torch.ones((edge_index.shape[1], 1)).to(device)

        emb, M = self.encoder(x, edge_index, edge_attr, batch, data.time)
        emb = self.proj_head(emb)
        
        return emb
    
    def loss_cal(self, pos, neg, T=0.1):
        pos = F.normalize(pos, dim=1)
        neg = F.normalize(neg, dim=1)
        sim_pp = torch.exp(pos @ pos.T / T)
        sim_pp = sim_pp - torch.diag_embed(torch.diag(sim_pp))

        sim_nn = torch.exp(neg @ neg.T / T)
        sim_nn = sim_nn - torch.diag_embed(torch.diag(sim_nn))

        sim_pn = torch.exp(pos @ neg.T / T)

        sum_pp = sim_pp.sum(1)
        sum_nn = sim_nn.sum(1)
        sum_pn = sim_pn.sum(1)

        loss_pos = -torch.log(sum_pp / (sum_pp+sum_pn)).mean()
        loss_neg = -torch.log(sum_nn / (sum_nn+sum_pn)).mean()
        return loss_pos + loss_neg