net_BNN.py

import math

import pickle

import torch
import torch.nn as nn
from torch.distributions import Normal
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.preprocessing import StandardScaler,MinMaxScaler
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F

class CycleDataset(Dataset):
    def __init__(self, data, attrib_x, attrib_y,  max_len, C_rated, min_val=None, max_val=None,  mode='train'):
        self.data = data
        self.cycle_indices = data['Cycle_Index'].unique()
        self.attrib_x = attrib_x
        self.attrib_y = attrib_y
        self.C_rated = C_rated
        self.mode = mode
        self.max_len = max_len

        self.data['Current'] /= self.C_rated
        if mode == 'train':
            self.min_val = data[attrib_x].values.min(axis=0)
            self.max_val = data[attrib_x].values.max(axis=0)
            with open('./para_BNN/min_max_values.pkl', 'wb') as f:
                pickle.dump((self.min_val, self.max_val), f)
        else:
            self.min_val = min_val
            self.max_val = max_val


    def get_min_max_values(self):
        if self.mode != 'train':
            return None
        return self.min_val, self.max_val


    def __len__(self):
        return len(self.cycle_indices)

    def __getitem__(self, index):
        cycle_index = self.cycle_indices[index]
        cycle_data = self.data[self.data['Cycle_Index'] == cycle_index].copy()

        # cycle_data['Current'] /= self.C_rated

        # 提取特征和标签
        features = cycle_data[self.attrib_x].values
        # C_ini = cycle_data[self.attrib_y].values[0]
        label = cycle_data[self.attrib_y].values[0]

        # # 标准化特征
        features = (features - self.min_val) / (self.max_val - self.min_val)
        # label = (label - self.y_mean) / self.y_std
        # features = (features - self.min_val) / self.max_val
        pad_len = self.max_len - len(features)

        features = torch.tensor(features, dtype=torch.float32).clone().detach()
        # 在 features 后面填充固定值
        features = torch.cat([features, torch.full((pad_len, features.shape[1]), 0)])
        # 转换为张量
        # features = torch.tensor(padded_features, dtype=torch.float32)
        label = torch.tensor(label, dtype=torch.float32)
        # label = label.view(1,1)

        return features, label

    #
    # def pad_collate(self, batch):
    #     # 填充批次数据，使其长度一致
    #     features_batch, labels_batch = zip(*batch)
    #     features_batch = pad_sequence(features_batch, batch_first=True)
    #     labels_batch = torch.stack(labels_batch)
    #
    #     return features_batch, labels_batch



class Transformer_FeatureExtractor(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, num_layers, num_heads, batch_size, max_seq_len):
        super(Transformer_FeatureExtractor, self).__init__()

        self.num_layers = num_layers
        self.hidden_size = hidden_dim
        self.batch_size = batch_size
        self.max_seq_len = max_seq_len
        # self.cls_token = nn.Parameter(torch.randn(self.batch_size, 1, self.hidden_size))
        self.embedding = nn.Linear(input_dim, hidden_dim)
        self.position_encoding = self.create_position_encoding()

        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=num_heads,dropout=0),
            num_layers=num_layers
        )

    def create_position_encoding(self):
        position_encoding = torch.zeros(self.max_seq_len, self.hidden_size)
        position = torch.arange(0, self.max_seq_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, self.hidden_size, 2) * (-math.log(10000.0) / self.hidden_size))
        position_encoding[:, 0::2] = torch.sin(position * div_term)
        position_encoding[:, 1::2] = torch.cos(position * div_term)
        position_encoding = position_encoding.unsqueeze(0)
        return nn.Parameter(position_encoding, requires_grad=False)

    def forward(self, x):
        seq_len = x.shape[1]
        positions = self.position_encoding[:, :seq_len, :]
        x = self.embedding(x)
        x = x + positions
        # x = torch.cat((x, self.cls_token),dim=1)
        # x = torch.cat((self.cls_token, x),dim=1)
        x_layer = self.transformer_encoder(x)
        feature = torch.mean(x_layer, dim=1)
        return feature

# class BaseVaraitionLayer_(nn.Module):
#     def __init__(self):
#         super().__init__()
#     def kl_div(self, mu_q, sigma_q, mu_p, sigma_p):
#         '''
#         Calculates kl divergence between two guassians (Q || P)
#         :param mu_q: torch.Tensor -> mu parameter of distribution Q
#         :param sigma_q: torch.Tensor -> sigma parameter of distribution Q
#         :param mu_p: float -> mu parameter of distribution P
#         :param sigma_p: float -> sigma parameter of distribution P
#         :return: torch.Tensor of shape 0
#         '''
#         kl = torch.log(sigma_p) - torch.log(sigma_q)
#         + (sigma_q**2 + (mu_q - mu_p)**2) / (2 * (sigma_p**2)) - 0.5
#         return kl.sum()


class BayesLinear(nn.Module):
    def __init__(self, input_dim, output_dim, prior_mu, prior_sigma):
        super(BayesLinear, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.prior_mu = prior_mu
        self.prior_sigma = prior_sigma

        self.weight_mu = nn.Parameter(torch.Tensor(output_dim, input_dim))
        self.weight_rho = nn.Parameter(torch.Tensor(output_dim, input_dim))
        self.bias_mu = nn.Parameter(torch.Tensor(output_dim))
        self.bias_rho = nn.Parameter(torch.Tensor(output_dim))

        self.weight = None
        self.bias = None

        self.prior = Normal(prior_mu, prior_sigma)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight_mu, a=math.sqrt(self.input_dim))
        nn.init.constant_(self.weight_rho, -3.0)
        nn.init.zeros_(self.bias_mu)
        nn.init.constant_(self.bias_rho, -3.0)

    def forward(self, input):
        weight_epsilon = torch.randn_like(self.weight_mu)
        bias_epsilon = torch.randn_like(self.bias_mu)

        weight_sigma = torch.log1p(torch.exp(self.weight_rho))
        bias_sigma = torch.log1p(torch.exp(self.bias_rho))

        self.weight = self.weight_mu + weight_sigma * weight_epsilon
        self.bias = self.bias_mu + bias_sigma * bias_epsilon

        weight_log_prior = self.prior.log_prob(self.weight)
        bias_log_prior = self.prior.log_prob(self.bias)
        self.log_prior = torch.sum(weight_log_prior) + torch.sum(bias_log_prior)

        self.weight_post = Normal(self.weight_mu.data, torch.log(1 + torch.exp(self.weight_rho)))
        self.bias_post = Normal(self.bias_mu.data, torch.log(1 + torch.exp(self.bias_rho)))
        self.log_post = self.weight_post.log_prob(self.weight).sum() + self.bias_post.log_prob(self.bias).sum()

        # output_mean = torch.matmul(input, weight.t()) + bias
        # output_var = torch.matmul(input, weight_sigma.t())**2 + bias_sigma**2
        # output_mean = nn.functional.linear(input, self.weight_mu, self.bias_mu)
        # output_variance = nn.functional.linear(input ** 2, weight_sigma ** 2, bias_sigma ** 2) + 1e-8
        # return output_mean, output_var
        return F.linear(input, self.weight, self.bias)

class BNN_Regression(nn.Module):
    def __init__(self, input_dim, output_dim, noise_tol):
        super(BNN_Regression, self).__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim
        # self.batch_size = batch_size
        self.noise_tol = noise_tol

        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
        # self.bnn1 = BayesLinear(input_dim=input_dim, output_dim=64, prior_mu=0, prior_sigma=1.)
        # self.bnn2 = BayesLinear(input_dim=64, output_dim=32, prior_mu=0, prior_sigma=1.)
        # self.fc = BayesLinear(input_dim=16, output_dim=output_dim,prior_mu=0, prior_sigma=1.)
        self.bnn = BayesLinear(input_dim=input_dim, output_dim=16, prior_mu=0, prior_sigma=1.)

        self.fc = BayesLinear(input_dim=16, output_dim=output_dim, prior_mu=0, prior_sigma=1.)
   
    def forward(self, x):
        x = self.bnn(x)
        x = self.relu(x)
        predictions = self.fc(x)
        # x = self.bnn1(x)
        # x = self.relu(x)
        # x = self.bnn2(x)
        # x = self.tanh(x)
        # x = self.bnn3(x)
        # x = self.relu(x)
        # predictions = self.fc(x)

        return predictions


    def log_prior(self):
        # calculate the log prior over all the layers
        # return self.bnn1.log_prior + self.bnn2.log_prior + self.bnn3.log_prior + self.fc.log_prior

        return self.bnn.log_prior + self.fc.log_prior

    def log_post(self):
        # calculate the log posterior over all the layers
        # return self.bnn1.log_post + self.bnn2.log_post + self.bnn3.log_post + self.fc.log_post

        return self.bnn.log_post + self.fc.log_post


    def sample_elbo(self, input, target, samples, device):
        # we calculate the negative elbo, which will be our loss function
        # initialize tensors
        outputs = torch.zeros(samples, target.shape[0]).to(device)
        log_priors = torch.zeros(samples)
        log_posts = torch.zeros(samples)
        log_likes = torch.zeros(samples)
        # make predictions and calculate prior, posterior, and likelihood for a given number of samples
        # 蒙特卡罗近似
        for i in range(samples):
            outputs[i] = self(input).reshape(-1)  # make predictions
            log_priors[i] = self.log_prior()  # get log prior
            log_posts[i] = self.log_post()  # get log variational posterior
            log_likes[i] = Normal(outputs[i], self.noise_tol).log_prob(target.reshape(-1)).sum()  # calculate the log likelihood
        # calculate monte carlo estimate of prior posterior and likelihood
        log_prior = log_priors.mean()
        log_post = log_posts.mean()
        log_like = log_likes.mean()
        # calculate the negative elbo (which is our loss function)
        loss = log_post - log_prior - log_like
        return loss


class ATBNN_Model(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, num_layers, num_heads, batch_size, max_seq_len):
        super(ATBNN_Model, self).__init__()

        self.feature_extractor = Transformer_FeatureExtractor(input_dim=input_dim,
                                                              output_dim=hidden_dim,
                                                              hidden_dim=hidden_dim,
                                                              num_layers=num_layers,
                                                              num_heads=num_heads,
                                                              batch_size=batch_size,
                                                              max_seq_len=max_seq_len)

        self.bnn_regression = BNN_Regression(input_dim=hidden_dim,
                                             output_dim=output_dim,
                                             noise_tol=0.01)

    def forward(self, x):
        self.features = self.feature_extractor(x)
        predictions = self.bnn_regression(self.features)
        return predictions