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