У меня есть код

Установка необходимых зависимостей
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_Uni.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[:,:2]

# Рабочая папка для моделей и графиков
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': []
    }
    
    cat_index = 0
    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([cat_index])
            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(cat_index))

        # 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([cat_index])
            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(cat_index))

        # 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([cat_index])
            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(cat_index))
            
        cat_index += 1

    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=[len(train_dataset)],
    # для каждого категорильного признака задаем размерность ембедингов:
    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) 

он использует класс from transformers import TimeSeriesTransformerConfig, TimeSeriesTransformerForPrediction

адаптируй данный код внеся изменения в TimeSeriesTransformerConfig, TimeSeriesTransformerForPrediction

так чтобы distribution_output стал Gamma, так как по умолчанию данный класс данное распределение не поддерживает, соответственно как внесешь изменения в класс, внеси изменения в процесс обучения, eval (так как для gamma распределения будут использоваться другие выходы модели на этапе обучения)