Продолжи код, я создаю модель трансформера для time series, который будет предсказывать распределение gamma вместо student_t

Вот мой старый код

Установка необходимых зависимостей
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) 


*** Здесь начинается новый код для модели кастомной модели которая будет предсказывать gamma распределение

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
from gluonts.torch.distributions import DistributionOutput
from gluonts.torch.distributions import GammaOutput

class CustomGammaOutput(DistributionOutput):
    @property
    def args_dim(self):
        return {'concentration': 1, 'rate': 1}

    @property
    def distribution_class(self):
        return GammaOutput
    
    def distribution(self, distr_args):
        conc, rate = distr_args
        return GammaOutput(conc, rate)
    
    def lower_bound_parameter(self):
        return 1e-8
config = {
    'prediction_length': prediction_length,
    'context_length': context_length,
    'num_feat_dynamic_real': len(time_features),
    'num_feat_static_real': 0,
    'num_feat_static_cat': 1,
    'cardinality': [len(train_dataset)],
    'embedding_dimension': [2],
    'dim_feedforward': 64,
    'model_dim': 32,
    'num_heads': 4,
    'num_encoder_layers': 4,
    'num_decoder_layers': 4,
    'dropout_rate': 0.1,
    'distribution_output': CustomGammaOutput(),
}
### Импорты и настройка модели:

import torch
import torch.nn as nn
import torch.optim as optim
from gluonts.dataset.common import ListDataset
from gluonts.evaluation import Evaluator
from gluonts.time_feature import get_lags_for_frequency, time_features_from_frequency_str
from gluonts.torch.model.deepar import DeepAREstimator
from gluonts.torch.distributions import DistributionOutput
from gluonts.torch.distributions import GammaOutput
from gluonts.torch.util import copy_parameters
from gluonts.dataset.loader import TrainDataLoader, InferenceDataLoader
from sklearn.metrics import mean_squared_error
from tqdm.auto import tqdm
import mlflow
import numpy as np


class CustomTransformer(nn.Module):
    def init(self, config):
        super(CustomTransformer, self).init()
        # Эти параметры следует задать в конфиге
        self.prediction_length = config['prediction_length']
        self.context_length = config['context_length']
        self.lags_sequence = config['lags_sequence']
        self.num_time_features = config['num_time_features']
        self.num_static_categorical_features = config['num_static_categorical_features']
        self.cardinality = config['cardinality']
        self.embedding_dimension = config['embedding_dimension']

        # Модель Transformer
        self.transformer = nn.Transformer(
            d_model=config['d_model'],
            nhead=config['num_heads'],
            num_encoder_layers=config['num_encoder_layers'],
            num_decoder_layers=config['num_decoder_layers'],
            dim_feedforward=config['dim_feedforward'],
            dropout=config['dropout_rate']
        )

        # GammaOutput для предсказания gamma распределения
        self.distribution_output = GammaOutput()

    def forward(
        self,
        past_time_features,
        past_values,
        future_time_features,
        past_observed_mask,
        future_values=None,
        static_categorical_features=None,
        static_real_features=None,
        future_observed_mask=None,
    ):
        # Реализация forward pass

        # Пример преобразований
        x = torch.cat((past_time_features, future_time_features), dim=1)
        output = self.transformer(x)

        # Пример с использованием GammaOutput
        distr_args = self.distribution_output.get_distr_args(output)
        distr = self.distribution_output.distribution(distr_args)
        
        loss = (-distr.log_prob(future_values) * future_observed_mask).sum() / future_observed_mask.sum()

        return {'loss': loss, 'distr_args': distr_args, 'output': output}

    def generate(
        self,
        past_time_features,
        past_values,
        future_time_features,
        past_observed_mask,
        static_categorical_features=None,
        static_real_features=None,
    ):
        # Реализация генерации предсказаний
        output = self.forward(
            past_time_features=past_time_features,
            past_values=past_values,
            future_time_features=future_time_features,
            past_observed_mask=past_observed_mask,
            static_categorical_features=static_categorical_features,
            static_real_features=static_real_features,
        )
        
        distr_args = output['distr_args']
        distr = self.distribution_output.distribution(distr_args)
        samples = distr.sample((self.num_samples,))
        
        return samples

Продолжи код, чтобы он был аналогичен моему базовому коду, только использовал кастомную модель трансформера