搭建GAT网络实现节点回归预测

猛男的成长之路

已于 2023-10-10 10:19:06 修改

阅读量921

点赞数 2

文章标签：回归数据挖掘人工智能

于 2023-10-09 22:59:52 首次发布

本文链接：https://blog.csdn.net/m0_57144920/article/details/133718662

版权

导入数据读取与可视化库

import pandas as pd
pd.set_option('display.max_columns',90)
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文字体设置-黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
sns.set(font='SimHei')  # 解决Seaborn中文显示问题并调整字体大小
#去除警告
import warnings
warnings.filterwarnings("ignore")

读数据

data = pd.read_csv('./相关性分析前所有特征数据.csv', encoding='utf_8_sig')
# data = pd.read_csv('./剔除特征高相关数据.csv', encoding='utf_8_sig')
data.head(5)

数据归一化与划分数据集

Input1, Output1 = data.drop(['barCode', '补偿容量'],axis=1).values, data['补偿容量'].values.reshape(-1, 1)

from sklearn.preprocessing import MinMaxScaler, StandardScaler
inputScaler = StandardScaler()
inputScaler = inputScaler.fit(Input1)
X = inputScaler.transform(Input1)

outputScaler = StandardScaler()
outputScaler = outputScaler.fit(Output1)
Y = outputScaler.transform(Output1)

from sklearn.model_selection import train_test_split
indices = np.arange(Input1.shape[0]) # help for check the index after split
[indices_train, indices_test, _, _] = train_test_split(indices, indices, test_size=0.2, random_state=1000, shuffle=True)
print(indices_train.shape, indices_test.shape)

train_x1, train_y1 = X[indices_train], Y[indices_train]
test_x1, test_y1 = X[indices_test], Y[indices_test]
train_y_org, test_y_org = Output1[indices_train], Output1[indices_test]
train_x1.shape, train_y1.shape, test_x1.shape, test_y1.shape

模型架构如下

import torch
import torch.nn.functional as F
import torch.nn as nn
from torch_geometric.nn.conv import GATConv

class GAT(nn.Module):
    def __init__(self, num_node_features):
        super(GAT, self).__init__()
        self.conv1 = GATConv(in_channels=num_node_features,out_channels=64,heads=4, dropout=0.5)
#         self.conv2 = GATConv(in_channels=2*256,out_channels=64,heads=1)
        self.layer2 = nn.Sequential(nn.Linear(64*4, 128), nn.BatchNorm1d(128), nn.ReLU(True), nn.Dropout(0.5))
#         self.resblock1 = nn.Sequential(nn.Linear(128, 128 // 2), nn.BatchNorm1d(128 // 2), nn.ReLU(True),
#                                      nn.Linear(128 // 2, 128 // 2), nn.BatchNorm1d(128 // 2), nn.ReLU(True),
#                                      nn.Linear(128 // 2, 128), nn.BatchNorm1d(128),  nn.ReLU(True))
        
        self.layer3 = nn.Sequential(nn.Linear(128, 64), nn.BatchNorm1d(64), nn.ReLU(True), nn.Dropout(0.5))
        self.out = nn.Linear(64, 1)
        
    def forward(self, x, edge_index):
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.layer2(x)
#         x = x + self.resblock1(x)
        x = self.layer3(x)
        x = self.out(x)
        return x

训练函数如下

from torch_cluster import knn_graph
import os
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import math

def trainer(model, config, device, train_x1, test_x1, train_y1, test_y1,train_org_y, test_org_y, outscaler):
    
    log = {'train_epochs_mse':[], 'test_epochs_mse':[]}
    result = {'train_r2':[], 'train_mse':[], 'train_mae':[], 'test_r2':[], 'test_mse':[], 'test_mae':[]}
    
    criterion = nn.MSELoss(reduction='mean')  # Define your loss function
#     criterion = nn.L1Loss(reduction='mean')

    optimizer = torch.optim.Adam(model.parameters(), lr=config['learning_rate'],  weight_decay=1e-6)

    if not os.path.isdir('./models'):
        os.mkdir('./models')  # Create directory of saving models.

    n_epochs, best_loss, step, early_stop_count = config['n_epochs'], math.inf, 0, 0

    # Create dataset
    train_count = train_x1.shape[0]
    X = torch.concatenate((train_x1, test_x1), axis=0)
    batch = torch.zeros(len(X))
    edge_index = knn_graph(X, k=config['k'], batch=batch, loop=False)
    X = X.to(device)   # Move your data to device.
    edge_index = edge_index.to(device)  # Move your data to device.
    train_y1, test_y1 = train_y1.to(device), test_y1.to(device)

    for epoch in range(n_epochs):
        model.train()  # Set your model to train mode.
        
        #batch training
        for i in range(train_count//512):
        
            optimizer.zero_grad()  # Set gradient to zero.
            pred = model(X, edge_index)
            loss = criterion(pred[i*512:(i+1)*512], train_y1[i*512:(i+1)*512])
            loss.backward()  # Compute gradient(backpropagation).
            optimizer.step()  # Update parameters.

#         optimizer.zero_grad()  # Set gradient to zero.
#         pred = model(X, edge_index)
#         loss = criterion(pred[:train_count], train_y1)
#         loss.backward()  # Compute gradient(backpropagation).
#         optimizer.step()  # Update parameters.
        
        
        mean_train_loss = loss.detach().item()
        log['train_epochs_mse'].append(mean_train_loss)

        model.eval()  # Set your model to evaluation mode.
        with torch.no_grad():
            pred = model(X, edge_index)
            loss = criterion(pred[train_count:], test_y1)
            mean_valid_loss = loss.item()

        log['test_epochs_mse'].append(mean_valid_loss)

        print(f'Epoch [{epoch + 1}/{n_epochs}]: Train loss: {mean_train_loss:.4f}, Valid loss: {mean_valid_loss:.4f}')

        if (epoch + 1) % 50 == 0:
            #eval model
            with torch.no_grad():
                pred = model(X, edge_index).cpu().numpy()
                pred = outscaler.inverse_transform(pred)
                train_pred, test_pred = pred[:train_count], pred[train_count:]
                train_r2, test_r2 = r2_score(train_org_y, train_pred), r2_score(test_org_y, test_pred)
                train_mse, test_mse = mean_squared_error(train_org_y, train_pred), mean_squared_error(test_org_y, test_pred)
                train_mae, test_mae = mean_absolute_error(train_org_y, train_pred), mean_absolute_error(test_org_y, test_pred)

            result['train_r2'].append(train_r2)
            result['train_mse'].append(train_mse)
            result['train_mae'].append(train_mae)
            result['test_r2'].append(test_r2)
            result['test_mse'].append(test_mse)
            result['test_mae'].append(test_mae)

            print(f'Epoch [{epoch + 1}/{n_epochs}]: train_r2: {train_r2:.4f}, train_mse: {train_mse:.4f}, train_mae: {train_mae:.4f}')
            print(f'Epoch [{epoch + 1}/{n_epochs}]: test_r2: {test_r2:.4f}, test_mse: {test_mse:.4f}, test_mae: {test_mae:.4f}\n')

        if mean_valid_loss < best_loss:
            best_loss = mean_valid_loss
            torch.save(model.state_dict(), config['save_path'])  # Save your best model
            print('Saving model with loss {:.3f}...'.format(best_loss))
            early_stop_count = 0
        else:
            early_stop_count += 1

        if early_stop_count >= config['early_stop']:
            print('\nModel is not improving, so we halt the training session.')
            return log, result
    return log, result

加载GPU和转换数据集

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
train_X, train_Y = torch.from_numpy(train_x1).float(), torch.from_numpy(train_y1).float()
test_X= torch.from_numpy(test_x1).float()
test_Y= torch.from_numpy(test_y1).float()

开始训练

config = {'n_epochs':3000, 'k':20,  'learning_rate':0.005, 'early_stop':300, 'save_path':'./models/dnn.pth'}
model = GAT(84).to(device)
log, result = trainer(model, config, device, train_X, test_X, train_Y, test_Y,train_y_org, test_y_org, outputScaler)

画图

plt.figure(figsize=(12,4))
plt.plot(log['train_epochs_mse'],'-o',label="train_loss")
plt.plot(log['test_epochs_mse'],'-o',label="valid_loss")
plt.title("epochs_loss")
plt.legend()
plt.show()