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