i have code with timeseries model using GluonTS for preparing data and timeseriestransformers for model

Установка необходимых зависимостей
import pandas as pd
import numpy as np
from evaluate import load
from gluonts.transform import InstanceSplitter
from sklearn.metrics import mean_absolute_error, mean_squared_error
from transformers import TimeSeriesTransformerConfig, TimeSeriesTransformerForPrediction
from datasets import Dataset, DatasetDict
from torch.optim import AdamW
from typing import Iterable
from gluonts.itertools import Cyclic, Cached
from gluonts.dataset.loader import as_stacked_batches
from gluonts.transform.sampler import InstanceSampler
from typing import Optional
from accelerate import Accelerator
from tqdm.auto import tqdm
from evaluate import load
from gluonts.time_feature import get_seasonality
import torch
import psutil
import tensorflow as tf
import os
import mlflow
import matplotlib.pyplot as plt
from functools import lru_cache
from functools import partial
from itertools import islice
from transformers import PretrainedConfig
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter
from matplotlib.backends.backend_pdf import PdfPages
from gluonts.time_feature import get_lags_for_frequency
from gluonts.time_feature import time_features_from_frequency_str
from gluonts.torch.distributions import GammaOutput
Проверка ресурсов
torch.cuda.is_available()
!nvidia-smi
#Сколько памяти в системе всего сейчас доступно
total_memory_bytes = psutil.virtual_memory().free
total_memory_gb = total_memory_bytes / (1024 * 1024 * 1024)
print(f"Free memory: {total_memory_gb:.2f} GB")
Загрузка датасетов
#Тренировочная выборка

path_uni = '//teconpc.local/ta/ТеконАвтоматизация/!ПРОЕКТЫ/СМД/Предиктивная диагностика/Файлы для сотрудников/!Отчеты Пронин И.С/1_Исследование методов улучшения качества прогноза (трансформеры)/Сигналы/07MAA10CT001/Train_Uni.xlsx'
path_uni = './signals/07MAA10CT001/Train_Multi (+Lables).xlsx'
path1_uni = os.path.normpath(path_uni)
df = pd.read_excel(path1_uni).iloc[:,:]

#Тестовая выборка (!!!Только один сигнал 07MAA10CT001)

path_test = '//teconpc.local/ta/ТеконАвтоматизация/!ПРОЕКТЫ/СМД/Предиктивная диагностика/Файлы для сотрудников/!Отчеты Пронин И.С/1_Исследование методов улучшения качества прогноза (трансформеры)/Сигналы/07MAA10CT001/Test_Anomaly.xlsx'
path_test = './signals/07MAA10CT001/Test_Anomaly.xlsx'
path1_test = os.path.normpath(path_test)
df_test = pd.read_excel(path1_test).iloc[:,:]

Рабочая папка для моделей и графиков
work_model_path = '//teconpc.local/ta/ТеконАвтоматизация/!ПРОЕКТЫ/СМД/Предиктивная диагностика/Файлы для сотрудников/!Отчеты Пронин И.С/1_Исследование методов улучшения качества прогноза (трансформеры)/models'
work_plots_path = '//teconpc.local/ta/ТеконАвтоматизация/!ПРОЕКТЫ/СМД/Предиктивная диагностика/Файлы для сотрудников/!Отчеты Пронин И.С/1_Исследование методов улучшения качества прогноза (трансформеры)/plots'

Предобработка
#Предобработка данных
def transform_dataframe(df:pd.DataFrame):
df['Время'] = pd.to_datetime(df['Время'])
df['hour'] = df['Время'].dt.hour
df['day_of_week'] = df['Время'].dt.dayofweek
df.set_index('Время', inplace=True)


return df
df = transform_dataframe(df)
df_test = transform_dataframe(df_test)
#Инициализация переменных для обучения

time_freq = '1min'
lag_freq = '10min'
context_length = 128
prediction_length = 64
batch_size = 64
step_in_dataset = 64
lags_sequence = get_lags_for_frequency(lag_freq)
time_features = time_features_from_frequency_str(time_freq)

Функция для создания наборов данных для датасетов
def create_datasets(df, context_length, prediction_length, step_in_dataset, isseparate=False, istest=False):
train_examples = {
'start': [],
'target': [],
'feat_static_cat': [],
'feat_dynamic_real': [],
'item_id': []
}
validation_examples = {
'start': [],
'target': [],
'feat_static_cat': [],
'feat_dynamic_real': [],
'item_id': []
}
test_examples = {
'start': [],
'target': [],
'feat_static_cat': [],
'feat_dynamic_real': [],
'item_id': []
}


step_size = 1

if isseparate:
    step_size = step_in_dataset
else:
    if istest:
        step_size = 1


for i in range(0, len(df) - context_length - 2 * prediction_length, step_size):
    # Training set
    train_target = df['07MAA10CT001'].iloc[i:i + context_length].tolist()
    if len(train_target) == context_length:
        train_examples['start'].append(df.index[i].isoformat())
        train_examples['target'].append(train_target)
        train_examples['feat_static_cat'].append([df.iloc[i, 4].item()])
        train_examples['feat_dynamic_real'].append([
            df['hour'].iloc[i:i + context_length].tolist(),
            df['day_of_week'].iloc[i:i + context_length].tolist()
        ])
        train_examples['item_id'].append(str(i))

    # Validation set
    val_target = df['07MAA10CT001'].iloc[i:i + context_length + prediction_length].tolist()
    if len(val_target) == context_length + prediction_length:
        validation_examples['start'].append(df.index[i].isoformat())
        validation_examples['target'].append(val_target)
        validation_examples['feat_static_cat'].append([df.iloc[i, 4].item()])
        validation_examples['feat_dynamic_real'].append([
            df['hour'].iloc[i:i + context_length + prediction_length].tolist(),
            df['day_of_week'].iloc[i:i + context_length + prediction_length].tolist()
        ])
        validation_examples['item_id'].append(str(i))

    # Test set
    test_target = df['07MAA10CT001'].iloc[i:i + context_length + 2 * prediction_length].tolist()
    if len(test_target) == context_length + 2 * prediction_length:
        test_examples['start'].append(df.index[i].isoformat())
        test_examples['target'].append(test_target)
        test_examples['feat_static_cat'].append([df.iloc[i, 4].item()])
        test_examples['feat_dynamic_real'].append([
            df['hour'].iloc[i:i + context_length + 2 * prediction_length].tolist(),
            df['day_of_week'].iloc[i:i + context_length + 2 * prediction_length].tolist()
        ])
        test_examples['item_id'].append(str(i))

return train_examples, validation_examples, test_examples
#Получаем датасеты из обучающей выборки
train_examples, validation_examples, test_examples = create_datasets(df, context_length, step_in_dataset, prediction_length, isseparate=True)

train_dataset = Dataset.from_dict(train_examples)
validation_dataset = Dataset.from_dict(validation_examples)
test_dataset = Dataset.from_dict(test_examples)

dataset_1 = DatasetDict({
'train': train_dataset,
'validation': validation_dataset,
'test': test_dataset
})
#Получаем датасеты из тестовой выборки
train_examples, validation_examples, test_examples = create_datasets(df_test, context_length, step_in_dataset, prediction_length, isseparate=True, istest=False)

train_dataset = Dataset.from_dict(train_examples)
validation_dataset = Dataset.from_dict(validation_examples)
test_dataset = Dataset.from_dict(test_examples)

dataset_2 = DatasetDict({
'train': train_dataset,
'validation': validation_dataset,
'test': test_dataset
})
#Контроль разделения train/val/test

train_example = dataset_1["train"][2]
validation_example = dataset_1["validation"][2]

test_example = dataset_2["test"][2]
figure, axes = plt.subplots()
axes.plot(train_example["target"], color="blue")
axes.plot(validation_example["target"], color="red", alpha=0.5)

axes.plot(test_example["target"], color="green", alpha=0.5)
plt.show()
train_dataset = dataset_1['train']
val_dataset = dataset_1['validation']
test_dataset = dataset_2['test']
@lru_cache(10_000)
def convert_to_pandas_period(date, freq):
return pd.Period(date, freq)

def transform_start_field(batch, freq):
batch["start"] = [convert_to_pandas_period(date, freq) for date in batch["start"]]
return batch
train_dataset.set_transform(partial(transform_start_field, freq=time_freq))
val_dataset.set_transform(partial(transform_start_field, freq=time_freq))
test_dataset.set_transform(partial(transform_start_field, freq=time_freq))
train_dataset, val_dataset, test_dataset
Инициализация модели
config = TimeSeriesTransformerConfig(
#длина прогноза
prediction_length=prediction_length,
# длина контекста:
context_length=context_length,
# лаг для расчета позиционного ембединга через get_lags_for_frequency:
lags_sequence=lags_sequence,
#количество временных признаков:
num_time_features=len(time_features) + 1,
# количество категориальных признаков:
num_static_categorical_features=1,
# размерность категориальных признаков, для бинарной классификации будет [2]:
cardinality=[2],
# для каждого категорильного признака задаем размерность ембедингов:
embedding_dimension=[2],
distribution_output='negative_binomial',


# параметры трансформера:
encoder_layers=4,
decoder_layers=4,
d_model=32,

# # в случае мультивариативности добавляем настройки на количество фичей:
# num_target_features=n,
# target_indices=[0],
# include_target_feature=True,
)

model = TimeSeriesTransformerForPrediction(config)

дефолтное распределение которое строит модель и на основании которого считается ошибка
model.config.distribution_output
#зависимости для трансформации данных
from gluonts.time_feature import (
time_features_from_frequency_str,
TimeFeature,
get_lags_for_frequency,
)
from gluonts.dataset.field_names import FieldName
from gluonts.transform import (
AddAgeFeature,
AddObservedValuesIndicator,
AddTimeFeatures,
AsNumpyArray,
Chain,
ExpectedNumInstanceSampler,
InstanceSplitter,
RemoveFields,
SelectFields,
SetField,
TestSplitSampler,
Transformation,
ValidationSplitSampler,
VstackFeatures,
RenameFields,
)
Трансформация данных под конфиг модели
def create_transformation(freq: str, config: PretrainedConfig) -> Transformation:
remove_field_names = []
if config.num_static_real_features == 0:
remove_field_names.append(FieldName.FEAT_STATIC_REAL)
if config.num_dynamic_real_features == 0:
remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)
if config.num_static_categorical_features == 0:
remove_field_names.append(FieldName.FEAT_STATIC_CAT)


return Chain(
    [RemoveFields(field_names=remove_field_names)]
    + (
        [
            AsNumpyArray(
                field=FieldName.FEAT_STATIC_CAT,
                expected_ndim=1,
                dtype=int,
            )
        ]
        if config.num_static_categorical_features > 0
        else []
    )
    + (
        [
            AsNumpyArray(
                field=FieldName.FEAT_STATIC_REAL,
                expected_ndim=1,
            )
        ]
        if config.num_static_real_features > 0
        else []
    )
    + [
        AsNumpyArray(
            field=FieldName.TARGET,
            expected_ndim=1 if config.input_size == 1 else 2,
        ),
        AddObservedValuesIndicator(
            target_field=FieldName.TARGET,
            output_field=FieldName.OBSERVED_VALUES,
        ),
        AddTimeFeatures(
            start_field=FieldName.START,
            target_field=FieldName.TARGET,
            output_field=FieldName.FEAT_TIME,
            time_features=time_features_from_frequency_str(freq),
            pred_length=config.prediction_length,
        ),
        AddAgeFeature(
            target_field=FieldName.TARGET,
            output_field=FieldName.FEAT_AGE,
            pred_length=config.prediction_length,
            log_scale=True,
        ),
        VstackFeatures(
            output_field=FieldName.FEAT_TIME,
            input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
            + (
                [FieldName.FEAT_DYNAMIC_REAL]
                if config.num_dynamic_real_features > 0
                else []
            ),
        ),
        RenameFields(
            mapping={
                FieldName.FEAT_STATIC_CAT: "static_categorical_features",
                FieldName.FEAT_STATIC_REAL: "static_real_features",
                FieldName.FEAT_TIME: "time_features",
                FieldName.TARGET: "values",
                FieldName.OBSERVED_VALUES: "observed_mask",
            }
        ),
    ]
)
def create_instance_splitter(
config: PretrainedConfig,
mode: str,
train_sampler: Optional[InstanceSampler] = None,
validation_sampler: Optional[InstanceSampler] = None,
) -> Transformation:
assert mode in ["train", "validation", "test"]


instance_sampler = {
    "train": train_sampler
    or ExpectedNumInstanceSampler(
        num_instances=1.0, min_future=config.prediction_length
    ),
    "validation": validation_sampler
    or ValidationSplitSampler(min_future=config.prediction_length),
    "test": TestSplitSampler(),
}[mode]

return InstanceSplitter(
    target_field="values",
    is_pad_field=FieldName.IS_PAD,
    start_field=FieldName.START,
    forecast_start_field=FieldName.FORECAST_START,
    instance_sampler=instance_sampler,
    past_length=config.context_length + max(config.lags_sequence),
    future_length=config.prediction_length,
    time_series_fields=["time_features", "observed_mask"],
)
def create_train_dataloader(
config: PretrainedConfig,
freq,
data,
batch_size: int,
num_batches_per_epoch: int,
shuffle_buffer_length: Optional[int] = None,
cache_data: bool = True,
**kwargs,
) -> Iterable:
PREDICTION_INPUT_NAMES = [
"past_time_features",
"past_values",
"past_observed_mask",
"future_time_features",
]
if config.num_static_categorical_features > 0:
PREDICTION_INPUT_NAMES.append("static_categorical_features")


if config.num_static_real_features > 0:
    PREDICTION_INPUT_NAMES.append("static_real_features")

TRAINING_INPUT_NAMES = PREDICTION_INPUT_NAMES + [
    "future_values",
    "future_observed_mask",
]

transformation = create_transformation(freq, config)
transformed_data = transformation.apply(data, is_train=True)
if cache_data:
    transformed_data = Cached(transformed_data)

instance_splitter = create_instance_splitter(config, "train")

stream = Cyclic(transformed_data).stream()
training_instances = instance_splitter.apply(stream)

return as_stacked_batches(
    training_instances,
    batch_size=batch_size,
    shuffle_buffer_length=shuffle_buffer_length,
    field_names=TRAINING_INPUT_NAMES,
    output_type=torch.tensor,
    num_batches_per_epoch=num_batches_per_epoch,
)
def create_backtest_dataloader(
config: PretrainedConfig,
freq,
data,
batch_size: int,
**kwargs,
):
PREDICTION_INPUT_NAMES = [
"past_time_features",
"past_values",
"past_observed_mask",
"future_time_features",
]
if config.num_static_categorical_features > 0:
PREDICTION_INPUT_NAMES.append("static_categorical_features")


if config.num_static_real_features > 0:
    PREDICTION_INPUT_NAMES.append("static_real_features")

transformation = create_transformation(freq, config)
transformed_data = transformation.apply(data)

instance_sampler = create_instance_splitter(config, "validation")

testing_instances = instance_sampler.apply(transformed_data, is_train=True)

return as_stacked_batches(
    testing_instances,
    batch_size=batch_size,
    output_type=torch.tensor,
    field_names=PREDICTION_INPUT_NAMES,
)
def create_test_dataloader(
config: PretrainedConfig,
freq,
data,
batch_size: int,
**kwargs,
):
PREDICTION_INPUT_NAMES = [
"past_time_features",
"past_values",
"past_observed_mask",
"future_time_features",
]
if config.num_static_categorical_features > 0:
PREDICTION_INPUT_NAMES.append("static_categorical_features")


if config.num_static_real_features > 0:
    PREDICTION_INPUT_NAMES.append("static_real_features")

transformation = create_transformation(freq, config)
transformed_data = transformation.apply(data, is_train=False)

instance_sampler = create_instance_splitter(config, "test")

testing_instances = instance_sampler.apply(transformed_data, is_train=False)

return as_stacked_batches(
    testing_instances,
    batch_size=batch_size,
    output_type=torch.tensor,
    field_names=PREDICTION_INPUT_NAMES,
)
train_dataloader = create_train_dataloader(
config=config,
freq=time_freq,
data=train_dataset,
batch_size=batch_size,
num_batches_per_epoch=100,
)

val_dataloader = create_backtest_dataloader(
config=config,
freq=time_freq,
data=val_dataset,
batch_size=batch_size,
)

test_dataloader = create_test_dataloader(
config=config,
freq=time_freq,
data=test_dataset,
batch_size=batch_size,
)
Инициализация акселератора, перенос модели на GPU
#Инициализация cuda, перенос модели на GPU

accelerator = Accelerator()
device = accelerator.device

print(device)

model.to(device)
Инициализация эксперимента и трекинга в MLflow
TRAINING=True

mlflow настройка эксперимента
mlflow.set_tracking_uri("http://od-eto:5002/")

experiment_name = "transformer_probability_enc_dec"

try:
object = mlflow.get_experiment_by_name(experiment_name)
if object is not None:
experiment_id = object.experiment_id
else:
experiment_id = mlflow.create_experiment(experiment_name)
except Exception:
print(f'Failed with getting experiment_id')

print(f"Experiment ID: {experiment_id}")
Обучение

Расчет метрик для критерия ранней остановы
def evaluate(model, dataloader, accelerator, num_batches):
model.eval()
batch_count = 0


forecasts = []

for batch in tqdm(dataloader):
    with torch.no_grad():
        outputs = model.generate(
            static_categorical_features=batch["static_categorical_features"].to(device)
            if config.num_static_categorical_features > 0
            else None,
            static_real_features=batch["static_real_features"].to(device)
            if config.num_static_real_features > 0
            else None,
            past_time_features=batch["past_time_features"].to(device),
            past_values=batch["past_values"].to(device),
            future_time_features=batch["future_time_features"].to(device),
            past_observed_mask=batch["past_observed_mask"].to(device),
        )
        forecasts.append(outputs.sequences.cpu().numpy())
        batch_count += 1
        
        if batch_count >= num_batches:
            break

forecasts = np.vstack(forecasts)
forecast_mean = np.mean(forecasts, axis=1)

rmse_metrics = []
last_index = num_batches*batch_size
count_idx = 0

for item_id, ts in tqdm(enumerate(val_dataset)):
    ground_truth = ts["target"][-prediction_length:]
    
    if len(forecast_mean[item_id]) != len(ground_truth):
        raise ValueError(f"Несоответствие по индексы: {item_id}: "
                         f"forecast len {len(forecast_mean[item_id])}, "
                         f"ground truth len {len(ground_truth)}")

    rmse = np.sqrt(mean_squared_error(ground_truth, forecast_mean[item_id]))
    rmse_metrics.append(rmse)
    
    count_idx += 1
    
    if count_idx >= last_index:
        break
    
return np.mean(rmse_metrics)
class CustomEarlyStoppingCallback:
def init(self, patience=5, verbose=False, min_delta=0.0):
self.patience = patience
self.verbose = verbose
self.min_delta = min_delta
self.counter = 0
self.min_validation_loss = np.inf
self.early_stop = False


def step(self, validation_loss):
    if validation_loss < self.min_validation_loss - self.min_delta:
        self.min_validation_loss = validation_loss
        self.counter = 0
    else:
        self.counter += 1
        if self.counter >= self.patience:
            self.early_stop = True
            if self.verbose:
                print(f"Критерий остановы {self.counter} на эпохе без улучшения.")
#Обучение модели без коллбека на раннюю останову
# Инициализация оптимизатора
optimizer = AdamW(model.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=1e-1)
num_epoch = 50
best_loss = np.inf
best_model_path = None
run_name = f'Prob_forecast datasetstep={step_in_dataset} lag={lag_freq} in={context_length} out={prediction_length} no callback'
model_path = f'{work_model_path}/probability_model_rev_3_2'
# Подготовка модели, оптимизатора и dataloader'ов с использованием accelerator
model, optimizer, train_dataloader = accelerator.prepare(
model,
optimizer,
train_dataloader,
)
# Обучение модели с эпохами и колбеком
if TRAINING:
with mlflow.start_run(experiment_id=experiment_id, run_name=run_name):
mlflow.log_param("num_epoch", num_epoch)
mlflow.log_param("batch_size", batch_size)
mlflow.log_param("model_architecture", str(model))
model.train()
for epoch in tqdm(range(1)):
for idx, batch in enumerate(train_dataloader):
optimizer.zero_grad()
outputs = model(
static_categorical_features=batch["static_categorical_features"].to(device)
if config.num_static_categorical_features > 0
else None,
static_real_features=batch["static_real_features"].to(device)
if config.num_static_real_features > 0
else None,
past_time_features=batch["past_time_features"].to(device),
past_values=batch["past_values"].to(device),
future_time_features=batch["future_time_features"].to(device),
future_values=batch["future_values"].to(device),
past_observed_mask=batch["past_observed_mask"].to(device),
future_observed_mask=batch["future_observed_mask"].to(device),
)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
if idx % 100 == 0:
print(loss.item())
mlflow.log_metric("Loss -LNN", loss, step=epoch)
best_model_path = os.path.join(os.getcwd(), model_path)
torch.save(model.state_dict(), best_model_path)
mlflow.end_run()
model
#Обучение модели с колбеком на раннюю останову

Инициализация оптимизатора
optimizer = AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.98), weight_decay=1e-2)
num_val_batches = 20
num_epoch = 500
run_name = f'Prob_forecast datasetstep={step_in_dataset} lag={lag_freq} in={context_length} out={prediction_length}'
best_val_loss = np.inf
best_model_path = None
model_path = f'{work_model_path}/probability_model_rev_3_2'

Подготовка модели, оптимизатора и dataloader'ов с использованием accelerator
model, optimizer, train_dataloader = accelerator.prepare(
model,
optimizer,
train_dataloader,
)

Инициализация колбека на раннюю остановку
early_stopping = CustomEarlyStoppingCallback(patience=10, min_delta=1e-4, verbose=True)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.5)

Обучение модели с эпохами и колбеком
if TRAINING:


with mlflow.start_run(experiment_id=experiment_id, run_name=run_name):
    mlflow.log_param("num_epoch", num_epoch)
    mlflow.log_param("batch_size", batch_size)
    mlflow.log_param("numbers_of_validation_batches", num_val_batches)
    mlflow.log_param("model_architecture", str(model))
    
    model.train()
    
    for epoch in tqdm(range(num_epoch)):
        for idx, batch in enumerate(train_dataloader):
            optimizer.zero_grad()
            outputs = model(
                static_categorical_features=batch["static_categorical_features"].to(device)
                if config.num_static_categorical_features > 0
                else None,
                static_real_features=batch["static_real_features"].to(device)
                if config.num_static_real_features > 0
                else None,
                past_time_features=batch["past_time_features"].to(device),
                past_values=batch["past_values"].to(device),
                future_time_features=batch["future_time_features"].to(device),
                future_values=batch["future_values"].to(device),
                past_observed_mask=batch["past_observed_mask"].to(device),
                future_observed_mask=batch["future_observed_mask"].to(device),
            )
            loss = outputs.loss

            accelerator.backward(loss)
            optimizer.step()

            if idx % 100 == 0:
                print(f"Epoch {epoch + 1}, Step {idx + 1}, Loss: {loss.item()}")


        val_loss_rmse = evaluate(model, val_dataloader, accelerator, num_val_batches)

        if val_loss_rmse < best_val_loss - early_stopping.min_delta:
            best_val_loss = val_loss_rmse
            best_model_path = os.path.join(os.getcwd(), model_path)
            torch.save(model.state_dict(), best_model_path)

        print(f"Epoch {epoch + 1}, Validation RMSE Loss: {val_loss_rmse}")
        
        mlflow.log_metric("Loss -LNN", loss, step=epoch)
        mlflow.log_metric("Val_loss RMSE", val_loss_rmse, step=epoch)

        scheduler.step(val_loss_rmse)
        early_stopping.step(val_loss_rmse)

        if early_stopping.early_stop:
            print("Сработали условия ранней остановы")
            break
            
    if best_model_path:
        model.load_state_dict(torch.load(best_model_path))
        print(f"Веса лучшей модели восстановлены: {best_model_path}")
    else:
        print("Нет лучшей сохраненной модели")
        
    mlflow.log_artifact(model_path)
    mlflow.end_run()

model
model_path = f'{work_model_path}/probability_model_rev_3_2'
model.load_state_dict(torch.load(model_path))
Оценка модели, формирование стака прогнозов, визуализация
EVALUATION = True
#Оценка модели с количеством батчей с начала датасета

if EVALUATION:
model.eval()
forecasts = []
count = 0


for batch in tqdm(test_dataloader):
    outputs = model.generate(
        static_categorical_features=batch["static_categorical_features"].to(device)
        if config.num_static_categorical_features > 0
        else None,
        static_real_features=batch["static_real_features"].to(device)
        if config.num_static_real_features > 0
        else None,
        past_time_features=batch["past_time_features"].to(device),
        past_values=batch["past_values"].to(device),
        future_time_features=batch["future_time_features"].to(device),
        past_observed_mask=batch["past_observed_mask"].to(device),
    )
    forecasts.append(outputs.sequences.cpu().numpy())
    count += 1
    
    if count >= 40:
        break
        
forecasts = np.vstack(forecasts)
forecast_mean = np.mean(forecasts, axis=1)
print(f'Размерность матрицы с прогнозами: {forecasts.shape}')
#Расчет MASE
def compute_mase(y_true, y_pred, y_train, seasonality=1):
n = len(y_train)
d = np.abs(np.diff(y_train, n=seasonality)).sum() / (n - seasonality)
errors = np.abs(y_true - y_pred)
return errors.mean() / d

forecast_mean = np.mean(forecasts, 1)

Расчет через самописную функцию
seasonality = 1
mase_metrics_1 = []
for item_id, ts in tqdm(enumerate(test_dataset)):
training_data = np.array(ts['target'][:-prediction_length*2])
ground_truth = np.array(ts['target'][-prediction_length:])
mase_metric = compute_mase(ground_truth, forecast_mean[item_id], training_data, seasonality=seasonality)
mase_metrics_1.append(mase_metric)

print(f'Метрика MASE (самописная) на тестовой выборке: {np.mean(mase_metrics_1)}')

freq = '5min'
mase_metric = load("evaluate-metric/mase")

forecast_mean = np.mean(forecasts, 1)

mase_metrics_2 = []

for item_id, ts in enumerate(test_dataset):
training_data = ts["target"][:-prediction_length*2]
ground_truth = ts["target"][-prediction_length:]
mase = mase_metric.compute(
predictions=forecast_mean[item_id],
references=np.array(ground_truth),
training=np.array(training_data),

periodicity=get_seasonality(freq),

)
mase_metrics_2.append(mase["mase"])
print(f'Метрика MASE (через evaluate) на тестовой выборке:: {np.mean(mase_metrics_2)}')
#Расчет RMSE
if EVALUATION:
try:
mlflow.set_tracking_uri("http://od-eto:5002/")
runs = mlflow.search_runs(experiment_ids = experiment_id, order_by=["start_time desc"], max_results=1)
mlflow.start_run(run_id = runs['run_id'][0])
except Exception as e:
print(f'Run уже запущен {e}')


rmse_metrics = []

for item_id, ts in tqdm(enumerate(test_dataset)):
    ground_truth = ts["target"][-prediction_length:]

    if len(forecast_mean[item_id]) != len(ground_truth):
        raise ValueError(f"Mismatch in length for item_id {item_id}: "
                         f"forecast length {len(forecast_mean[item_id])}, "
                         f"ground truth length {len(ground_truth)}")

    rmse = np.sqrt(mean_squared_error(ground_truth, forecast_mean[item_id]))
    rmse_metrics.append(rmse)

if rmse_metrics:
    mlflow.log_metric('RMSE на тестовой выборке', np.mean(rmse_metrics))

print(f'RMSE на тестовом датасете: {np.mean(rmse_metrics)}')
#Визуализация результатов модели на графике

def plot(ts_index):
fig, ax = plt.subplots(figsize=(30,10))


index = pd.period_range(
    start=test_dataset[ts_index][FieldName.START],
    periods=len(test_dataset[ts_index][FieldName.TARGET]),
    freq=time_freq,
).to_timestamp()
resampled_median_forecast = np.median(forecasts[ts_index], axis=0)
interpolated_median_forecast = savgol_filter(
resampled_median_forecast,
window_length=10,
polyorder=3,
mode="interp"
)
resampled_mean_forecast = np.mean(forecasts[ts_index], axis=0)
interpolated_mean_forecast = savgol_filter(
resampled_mean_forecast,
window_length=10,
polyorder=3,
mode="interp"
)

ax.plot(
    index[-2 * prediction_length :],
    test_dataset[ts_index]["target"][-2 * prediction_length :],
    label="реальные данные",
    color='black'
)
plt.plot(
index[-prediction_length:],
np.mean(forecasts[ts_index], axis=0),
label="средняя",
color='orange'
)

plt.plot(
    index[-prediction_length:],
    np.median(forecasts[ts_index], axis=0),
    label="медиана",
    color='green'
)
plt.plot(
index[-prediction_length:],
interpolated_median_forecast,
label="сглаженная медиана",
color='purple'
)
plt.plot(
index[-prediction_length:],
interpolated_mean_forecast,
label="сглаженная средняя",
color='red'
)

plt.fill_between(
    index[-prediction_length:],
    forecasts[ts_index].mean(0) - forecasts[ts_index].std(axis=0),
    forecasts[ts_index].mean(0) + forecasts[ts_index].std(axis=0),
    alpha=0.3,
    interpolate=True,
    label="+/- 1-std",
)

plt.legend()
fig.autofmt_xdate()
start_time = index[-2 * prediction_length]
end_time = index[-1]
plt.title(f"График отображает вывод за период с {start_time} до {end_time}")
plt.show()

return fig
#Загружаем альбом результатов тестов в MLflow
try:
mlflow.set_tracking_uri("http://od-eto:5002/")
runs = mlflow.search_runs(experiment_ids = experiment_id, order_by=["start_time desc"], max_results=1)
mlflow.start_run(run_id = runs['run_id'][0])
except Exception as e:
print(f'Run уже запущен {e}')

plot_path = f'{work_plots_path}/График для {run_name} на тестовой выборке.pdf'
plot_path = os.path.join(os.getcwd(), plot_path)

with PdfPages(plot_path) as pdf:
for i in range(0, forecasts.shape[0]):
fig = plot(i)
pdf.savefig(fig, bbox_inches='tight')
plt.close(fig)


mlflow.log_artifact(plot_path)
mlflow.end_run()

for i in range(0, forecasts.shape[0]):
plot(i)
Expand class from transformers import TimeSeriesTransformerConfig, TimeSeriesTransformerForPrediction to Custom class which will be able to use Gamma distribution as output, after check train, eval and test proccess and apapt it to new model