from typing import Dict import torchinfo import tqdm, math import numpy as np import torch from torch import nn, optim from torch.utils.data import DataLoader from ..utils.torch_utility import EarlyStoppingTorch, get_gpu from ..utils.dataset import ForecastDataset class LSTMModel(nn.Module): def __init__(self, window_size, feats, hidden_dim, pred_len, num_layers, batch_size, device) -> None: super().__init__() self.pred_len = pred_len self.batch_size = batch_size self.feats = feats self.device = device self.lstm_encoder = nn.LSTM(input_size=feats, hidden_size=hidden_dim, num_layers=num_layers, batch_first=True) self.lstm_decoder = nn.LSTM(input_size=feats, hidden_size=hidden_dim, num_layers=num_layers, batch_first=True) self.relu = nn.GELU() self.fc = nn.Linear(hidden_dim, feats) def forward(self, src): _, decoder_hidden = self.lstm_encoder(src) cur_batch = src.shape[0] decoder_input = torch.zeros(cur_batch, 1, self.feats).to(self.device) outputs = torch.zeros(self.pred_len, cur_batch, self.feats).to(self.device) for t in range(self.pred_len): decoder_output, decoder_hidden = self.lstm_decoder(decoder_input, decoder_hidden) decoder_output = self.relu(decoder_output) decoder_input = self.fc(decoder_output) outputs[t] = torch.squeeze(decoder_input, dim=-2) return outputs class LSTMAD(): def __init__(self, window_size=100, pred_len=1, batch_size=128, epochs=50, lr=0.0008, feats=1, hidden_dim=20, num_layer=2, validation_size=0.2): super().__init__() self.__anomaly_score = None cuda = True self.y_hats = None self.cuda = cuda self.device = get_gpu(self.cuda) self.window_size = window_size self.pred_len = pred_len self.batch_size = batch_size self.epochs = epochs self.feats = feats self.hidden_dim = hidden_dim self.num_layer = num_layer self.lr = lr self.validation_size = validation_size print('self.device: ', self.device) self.model = LSTMModel(self.window_size, feats, hidden_dim, self.pred_len, num_layer, batch_size=self.batch_size, device=self.device).to(self.device) self.optimizer = optim.Adam(self.model.parameters(), lr=lr) self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.75) self.loss = nn.MSELoss() self.save_path = None self.early_stopping = EarlyStoppingTorch(save_path=self.save_path, patience=3) self.mu = None self.sigma = None self.eps = 1e-10 def fit(self, data): tsTrain = data[:int((1-self.validation_size)*len(data))] tsValid = data[int((1-self.validation_size)*len(data)):] train_loader = DataLoader( ForecastDataset(tsTrain, window_size=self.window_size, pred_len=self.pred_len), batch_size=self.batch_size, shuffle=True) valid_loader = DataLoader( ForecastDataset(tsValid, window_size=self.window_size, pred_len=self.pred_len), batch_size=self.batch_size, shuffle=False) for epoch in range(1, self.epochs + 1): self.model.train(mode=True) avg_loss = 0 loop = tqdm.tqdm(enumerate(train_loader),total=len(train_loader),leave=True) for idx, (x, target) in loop: x, target = x.to(self.device), target.to(self.device) # print('x: ', x.shape) # (bs, win, feat) # print('target: ', target.shape) # # (bs, pred_len, feat) self.optimizer.zero_grad() output = self.model(x) # print('output: ', output.shape) # (pred_len, bs, feat) output = output.view(-1, self.feats*self.pred_len) target = target.view(-1, self.feats*self.pred_len) loss = self.loss(output, target) loss.backward() self.optimizer.step() avg_loss += loss.cpu().item() loop.set_description(f'Training Epoch [{epoch}/{self.epochs}]') loop.set_postfix(loss=loss.item(), avg_loss=avg_loss/(idx+1)) self.model.eval() scores = [] avg_loss = 0 loop = tqdm.tqdm(enumerate(valid_loader),total=len(valid_loader),leave=True) with torch.no_grad(): for idx, (x, target) in loop: x, target = x.to(self.device), target.to(self.device) output = self.model(x) output = output.view(-1, self.feats*self.pred_len) target = target.view(-1, self.feats*self.pred_len) loss = self.loss(output, target) avg_loss += loss.cpu().item() loop.set_description(f'Validation Epoch [{epoch}/{self.epochs}]') loop.set_postfix(loss=loss.item(), avg_loss=avg_loss/(idx+1)) mse = torch.sub(output, target).pow(2) scores.append(mse.cpu()) valid_loss = avg_loss/max(len(valid_loader), 1) self.scheduler.step() self.early_stopping(valid_loss, self.model) if self.early_stopping.early_stop or epoch == self.epochs - 1: # fitting Gaussian Distribution if len(scores) > 0: scores = torch.cat(scores, dim=0) self.mu = torch.mean(scores) self.sigma = torch.var(scores) print(self.mu.size(), self.sigma.size()) if self.early_stopping.early_stop: print(" Early stopping<<<") break def decision_function(self, data): test_loader = DataLoader( ForecastDataset(data, window_size=self.window_size, pred_len=self.pred_len), batch_size=self.batch_size, shuffle=False ) self.model.eval() scores = [] y_hats = [] loop = tqdm.tqdm(enumerate(test_loader),total=len(test_loader),leave=True) with torch.no_grad(): for idx, (x, target) in loop: x, target = x.to(self.device), target.to(self.device) output = self.model(x) output = output.view(-1, self.feats*self.pred_len) target = target.view(-1, self.feats*self.pred_len) mse = torch.sub(output, target).pow(2) y_hats.append(output.cpu()) scores.append(mse.cpu()) loop.set_description(f'Testing: ') scores = torch.cat(scores, dim=0) # scores = 0.5 * (torch.log(self.sigma + self.eps) + (scores - self.mu)**2 / (self.sigma+self.eps)) scores = scores.numpy() scores = np.mean(scores, axis=1) y_hats = torch.cat(y_hats, dim=0) y_hats = y_hats.numpy() l, w = y_hats.shape # new_scores = np.zeros((l - self.pred_len, w)) # for i in range(w): # new_scores[:, i] = scores[self.pred_len - i:l-i, i] # scores = np.mean(new_scores, axis=1) # scores = np.pad(scores, (0, self.pred_len - 1), 'constant', constant_values=(0,0)) # new_y_hats = np.zeros((l - self.pred_len, w)) # for i in range(w): # new_y_hats[:, i] = y_hats[self.pred_len - i:l-i, i] # y_hats = np.mean(new_y_hats, axis=1) # y_hats = np.pad(y_hats, (0, self.pred_len - 1), 'constant',constant_values=(0,0)) assert scores.ndim == 1 # self.y_hats = y_hats # print('scores: ', scores.shape) if scores.shape[0] < len(data): padded_decision_scores_ = np.zeros(len(data)) padded_decision_scores_[: self.window_size+self.pred_len-1] = scores[0] padded_decision_scores_[self.window_size+self.pred_len-1 : ] = scores self.__anomaly_score = padded_decision_scores_ return padded_decision_scores_ def anomaly_score(self) -> np.ndarray: return self.__anomaly_score def get_y_hat(self) -> np.ndarray: return self.y_hats def param_statistic(self, save_file): model_stats = torchinfo.summary(self.model, (self.batch_size, self.window_size), verbose=0) with open(save_file, 'w') as f: f.write(str(model_stats))