Spaces:
Running
Running
| """ | |
| This function is adapted from [moment] by [mononitogoswami] | |
| Original source: [https://github.com/moment-timeseries-foundation-model/moment] | |
| """ | |
| from momentfm import MOMENTPipeline | |
| from momentfm.utils.masking import Masking | |
| from momentfm.utils.utils import control_randomness | |
| from sklearn.preprocessing import MinMaxScaler | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from tqdm import tqdm | |
| from torch import nn | |
| import math | |
| from .base import BaseDetector | |
| from ..utils.dataset import ReconstructDataset_Moment | |
| from ..utils.torch_utility import EarlyStoppingTorch, get_gpu | |
| class MOMENT(BaseDetector): | |
| def __init__(self, | |
| win_size=256, | |
| input_c=1, | |
| batch_size=128, | |
| epochs=2, | |
| validation_size=0, | |
| lr=1e-4): | |
| self.model_name = 'MOMENT' | |
| self.win_size = win_size | |
| self.input_c = input_c | |
| self.batch_size = batch_size | |
| self.anomaly_criterion = nn.MSELoss(reduce=False) | |
| self.epochs = epochs | |
| self.validation_size = validation_size | |
| self.lr = lr | |
| cuda = True | |
| self.cuda = cuda | |
| self.device = get_gpu(self.cuda) | |
| # Control randomness for reproducibility | |
| control_randomness(seed=42) | |
| # Load the model properly with config | |
| try: | |
| self.model = MOMENTPipeline.from_pretrained( | |
| "AutonLab/MOMENT-1-base", | |
| model_kwargs={ | |
| "task_name": "reconstruction", | |
| "n_channels": self.input_c, | |
| "max_seq_len": self.win_size | |
| } | |
| ) | |
| self.model.init() | |
| except Exception as e: | |
| # Fallback: try alternative initialization | |
| print(f"Failed to load MOMENT model with from_pretrained: {e}") | |
| print("Attempting alternative initialization...") | |
| from transformers import AutoConfig | |
| config = AutoConfig.from_pretrained("AutonLab/MOMENT-1-base") | |
| config.task_name = "reconstruction" | |
| config.n_channels = self.input_c | |
| config.max_seq_len = self.win_size | |
| self.model = MOMENTPipeline(config) | |
| self.model.init() | |
| self.model = self.model.to(self.device).float() | |
| # Optimize Mean Squarred Error using your favourite optimizer | |
| self.criterion = torch.nn.MSELoss() | |
| self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr) | |
| self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.75) | |
| self.save_path = None | |
| self.early_stopping = EarlyStoppingTorch(save_path=self.save_path, patience=3) | |
| def zero_shot(self, data): | |
| test_loader = DataLoader( | |
| dataset=ReconstructDataset_Moment(data, window_size=self.win_size, normalize=True), | |
| batch_size=self.batch_size, | |
| shuffle=False) | |
| trues, preds = [], [] | |
| self.score_list = [] | |
| with torch.no_grad(): | |
| for batch_x, batch_masks in tqdm(test_loader, total=len(test_loader)): | |
| batch_x = batch_x.to("cuda").float() | |
| batch_masks = batch_masks.to("cuda") | |
| batch_x = batch_x.permute(0,2,1) | |
| # print('batch_x: ', batch_x.shape) # [batch_size, n_channels, window_size] | |
| # print('batch_masks: ', batch_masks.shape) # [batch_size, window_size] | |
| output = self.model(x_enc=batch_x, input_mask=batch_masks) # [batch_size, n_channels, window_size] | |
| score = torch.mean(self.anomaly_criterion(batch_x, output.reconstruction), dim=-1).detach().cpu().numpy()[:, -1] | |
| self.score_list.append(score) | |
| self.__anomaly_score = np.concatenate(self.score_list, axis=0).reshape(-1) | |
| if self.__anomaly_score.shape[0] < len(data): | |
| self.__anomaly_score = np.array([self.__anomaly_score[0]]*math.ceil((self.win_size-1)/2) + | |
| list(self.__anomaly_score) + [self.__anomaly_score[-1]]*((self.win_size-1)//2)) | |
| self.decision_scores_ = self.__anomaly_score | |
| 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( | |
| dataset=ReconstructDataset_Moment(tsTrain, window_size=self.win_size), | |
| batch_size=self.batch_size, | |
| shuffle=True | |
| ) | |
| valid_loader = DataLoader( | |
| dataset=ReconstructDataset_Moment(tsValid, window_size=self.win_size), | |
| batch_size=self.batch_size, | |
| shuffle=False | |
| ) | |
| mask_generator = Masking(mask_ratio=0.3) # Mask 30% of patches randomly | |
| for epoch in range(1, self.epochs + 1): | |
| self.model.train() | |
| for batch_x, batch_masks in tqdm(train_loader, total=len(train_loader)): | |
| batch_x = batch_x.to(self.device).float() | |
| batch_x = batch_x.permute(0,2,1) | |
| # print('batch_x: ', batch_x.shape) | |
| original = batch_x | |
| n_channels = batch_x.shape[1] | |
| # Reshape to [batch_size * n_channels, 1, window_size] | |
| batch_x = batch_x.reshape((-1, 1, self.win_size)) | |
| batch_masks = batch_masks.to(self.device).long() | |
| batch_masks = batch_masks.repeat_interleave(n_channels, axis=0) | |
| # Randomly mask some patches of data | |
| mask = mask_generator.generate_mask( | |
| x=batch_x, input_mask=batch_masks).to(self.device).long() | |
| mask = torch.nn.functional.pad(mask, (0, batch_masks.size(1) - mask.size(1)), mode='constant', value=1) | |
| # Forward | |
| model_output = self.model(batch_x, input_mask=batch_masks, mask=mask).reconstruction | |
| model_output = torch.nn.functional.pad(model_output, (0, original.size(2)-model_output.size(2)), mode='replicate') | |
| output = model_output.reshape(original.size(0), n_channels, self.win_size) | |
| # Compute loss | |
| loss = self.criterion(output, original) | |
| # print(f"loss: {loss.item()}") | |
| # Backward | |
| self.optimizer.zero_grad() | |
| loss.backward() | |
| self.optimizer.step() | |
| # self.model.eval() | |
| # avg_loss = 0 | |
| # with torch.no_grad(): | |
| # for batch_x, batch_masks in tqdm(valid_loader, total=len(valid_loader)): | |
| # batch_x = batch_x.to("cuda").float() | |
| # batch_masks = batch_masks.to("cuda") | |
| # batch_x = batch_x.permute(0,2,1) | |
| # print('batch_x: ', batch_x.shape) | |
| # print('batch_masks: ', batch_masks.shape) | |
| # output = self.model(batch_x, input_mask=batch_masks) | |
| # loss = self.criterion(output.reconstruction.reshape(-1, n_channels, self.win_size), batch_x) | |
| # print(f"loss: {loss.item()}") | |
| # avg_loss += loss.cpu().item() | |
| # 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: | |
| # print(" Early stopping<<<") | |
| # break | |
| def decision_function(self, data): | |
| """ | |
| Not used, present for API consistency by convention. | |
| """ | |
| test_loader = DataLoader( | |
| dataset=ReconstructDataset_Moment(data, window_size=self.win_size), | |
| batch_size=self.batch_size, | |
| shuffle=False) | |
| trues, preds = [], [] | |
| self.score_list = [] | |
| with torch.no_grad(): | |
| for batch_x, batch_masks in tqdm(test_loader, total=len(test_loader)): | |
| batch_x = batch_x.to("cuda").float() | |
| batch_masks = batch_masks.to("cuda") | |
| batch_x = batch_x.permute(0,2,1) | |
| # print('batch_x: ', batch_x.shape) # [batch_size, n_channels, window_size] | |
| # print('batch_masks: ', batch_masks.shape) # [batch_size, window_size] | |
| output = self.model(batch_x, input_mask=batch_masks) | |
| score = torch.mean(self.anomaly_criterion(batch_x, output.reconstruction), dim=-1).detach().cpu().numpy()[:, -1] | |
| self.score_list.append(score) | |
| self.__anomaly_score = np.concatenate(self.score_list, axis=0).reshape(-1) | |
| if self.__anomaly_score.shape[0] < len(data): | |
| self.__anomaly_score = np.array([self.__anomaly_score[0]]*math.ceil((self.win_size-1)/2) + | |
| list(self.__anomaly_score) + [self.__anomaly_score[-1]]*((self.win_size-1)//2)) | |
| return self.__anomaly_score |