softordering.py - master - TolmachevM/MLLowLevelResearch

softordering.py
437 строк26 KB
TolmachevM
update softordering.py
05 авг 2025, 06:45
05 авг 2025, 06:454de915c
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_sequence

# Использование, создание слоя:
# ordering = SoftOrdering(n_bins=5, gcoef=6.0, reduction='mean', input_sigmoid=True)
# Слой SoftOrdering() размещается после линейного слоя Linear() вместо функции активации.
# Слой SoftOrdering2d() размещается после свёрточного слоя Conv2d() вместо функции активации.
# При n_bins = 1 работает как бинарный слой мягкого упорядочивания BinarySoftOrdering().
# Если input_sigmoid=False, то код работает чуть быстрее,
# но такие слои могут применяются только к нормализованным в диапазоне (0.0, 1.0) данным.

class SoftOrdering(nn.Module):
    """
    Слой мягкого упорядочивания для линейных слоёв Linear().
    """
    def __init__(self, n_bins: int, gcoef: float = 6.0, reduction: str = 'mean', input_sigmoid: bool = True) -> None:
        super().__init__()
        self.n_bins = n_bins
        self.gcoef = gcoef
        self.input_sigmoid = input_sigmoid
        self.reduction = reduction
        self.sigmoid = nn.Sigmoid()
        if input_sigmoid:
            self.get_preprocessed_input = self._get_preprocessed_input_sigmoid
        else:
            self.get_preprocessed_input = self._get_preprocessed_input_pass
        self.sigmoid = nn.Sigmoid()
        if n_bins >= 2:
            self.gain = gcoef * (n_bins - 1)
            #self.gain = nn.Parameter(torch.Tensor([gcoef * (n_bins - 1)]), requires_grad=True)
            self.register_buffer('lspace', torch.linspace(-0.1 - 1/(n_bins - 1), 1.1 + 1/(n_bins - 1), steps=n_bins + 1))
            self.forward = self._forward_nonbinary
        else:
            self.forward = self._forward_binary
        if reduction == 'none':
            self._reduce = self._reduction_pass
        elif reduction == 'mean':
            self._reduce = self._reduction_mean
        else:
            assert "Unknown reduction function. Use 'none' or 'mean'."
    def extra_repr(self) -> str:
        reduction_info = "reduction=mean" if self.reduction == 'reduction=mean (Soft Ordering)' else "reduction=none (Soft One-Hot Encoding)"
        sigmoid_info = "input_sigmoid=True" if self.input_sigmoid else "input_sigmoid=False"
        return f"n_bins={self.n_bins}, gcoef={self.gcoef}, {reduction_info}, {sigmoid_info}"
    def _get_preprocessed_input_pass(self, x: torch.Tensor) -> torch.Tensor:
        return x
    def _get_preprocessed_input_sigmoid(self, x: torch.Tensor) -> torch.Tensor:
        return self.sigmoid(x)
    def _forward_nonbinary(self, x: torch.Tensor) -> torch.Tensor:
        # размерность тензора x произвольная, например, (batch, points, features),
        # последняя размерность features - обязательна,
        # предпоследняя размерность points обязательна для reduction='mean', необязательна для reduction='none'
        z = self.get_preprocessed_input(x)
        z = z.unsqueeze(-2)
        *batch_dims, n_points, n_features = x.shape # размерности        
        z_expanded = z.expand(*batch_dims, n_points, self.n_bins + 1, n_features) # Используем expand вместо repeat для экономии памяти
        lspace_shape = [1] * len(batch_dims) + [1, -1, 1] # Создаем lspace с правильными размерностями
        lspace_expanded = self.lspace.view(*lspace_shape)
        lspace_expanded = lspace_expanded.expand(*batch_dims, n_points, self.n_bins + 1, n_features)
        diff = z_expanded - lspace_expanded # Вычисляем разность и применяем sigmoid с gain
        z_sigmoid = self.sigmoid(self.gain * diff)
        # Более эффективное вычисление разности соседних элементов
        z_diff = z_sigmoid[..., :-1, :] - z_sigmoid[..., 1:, :] #размерность z_diff (batch, points, bins, features)
        z_flat = z_diff.flatten(-2, -1) # Flatten последних двух измерений
        # размерность тензора p повторяет первые размерности входного тензора x, например, (batch, points, bins * features)
        # в случае reduction='mean' размерность point сокращается, т.е. (batch, bins * features)
        return self._reduce(z_flat)
    def _forward_binary(self, x: torch.Tensor) -> torch.Tensor: #BinarySoftOrdering
        return self._reduce(self.sigmoid(x))
    def _reduction_pass(self, x: torch.Tensor) -> torch.Tensor:
        return x
    def _reduction_mean(self, x: torch.Tensor) -> torch.Tensor:
        return x.mean(-2)

        
class SoftOrdering2d(nn.Module):
    '''
    Ленивый слой мягкого упорядочивания для свёрточных слоёв Conv2d()
    '''
    def __init__(self, n_bins: int, gcoef: float = 6.0, reduction: str = 'mean', input_sigmoid: bool = True) -> None:
        super().__init__()
        self.n_bins = n_bins
        self.gcoef = gcoef
        self.input_sigmoid = input_sigmoid
        self.reduction = reduction
        self.sigmoid = nn.Sigmoid()
        if input_sigmoid:
            self.get_preprocessed_input = self._get_preprocessed_input_sigmoid
        else:
            self.get_preprocessed_input = self._get_preprocessed_input_pass
        if n_bins >= 2:
            self.gain = gcoef * (n_bins - 1)
            #self.gain = nn.Parameter(torch.Tensor([gcoef * (n_bins - 1)]), requires_grad=True)
            self.register_buffer('lspace', torch.linspace(-0.1 - 1/(n_bins - 1), 1.1 + 1/(n_bins - 1), steps=n_bins + 1))
            self.forward = self._forward_nonbinary
        else:
            self.forward = self._forward_binary
        if reduction == 'none':
            self._reduce = self._reduction_pass
        elif reduction == 'mean':
            self._reduce = self._reduction_mean
        else:
            assert "Unknown reduction function. Use 'none' or 'mean'."
    def extra_repr(self) -> str:
        reduction_info = "reduction=mean" if self.reduction == 'reduction=mean (Soft Ordering)' else "reduction=none (Soft One-Hot Encoding)"
        sigmoid_info = "input_sigmoid=True" if self.input_sigmoid else "input_sigmoid=False"
        return f"n_bins={self.n_bins}, gcoef={self.gcoef}, {reduction_info}, {sigmoid_info}"
    def _get_preprocessed_input_pass(self, x: torch.Tensor) -> torch.Tensor:
        return x
    def _get_preprocessed_input_sigmoid(self, x: torch.Tensor) -> torch.Tensor:
        return self.sigmoid(x)
    def _forward_nonbinary(self, x: torch.Tensor) -> torch.Tensor:
        # размерность тензора x произвольная, например, (batch, height, width, features),
        # последние размерности features - обязательны,
        # предпоследняя размерность points обязательна для reduction='mean', необязательна для reduction='none'
        z = self.get_preprocessed_input(x)
        z = z.unsqueeze(-2)
        *batch_dims, height, width, n_features = x.shape # размерности
        z_expanded = z.expand(*batch_dims, height, width, self.n_bins + 1, n_features) # Используем expand вместо repeat для экономии памяти
        lspace_shape = [1] * len(batch_dims) + [1, 1, -1, 1] # Создаем lspace с правильными размерностями
        lspace_expanded = self.lspace.view(*lspace_shape)
        lspace_expanded = lspace_expanded.expand(*batch_dims, height, width, self.n_bins + 1, n_features)
        diff = z_expanded - lspace_expanded # Вычисляем разность и применяем sigmoid с gain
        z_sigmoid = self.sigmoid(self.gain * diff)
        # Более эффективное вычисление разности соседних элементов
        z_diff = z_sigmoid[..., :-1, :] - z_sigmoid[..., 1:, :] #размерность z_diff (batch, points, bins, features)
        z_moved = z_diff.movedim(-2, -3)  # (batch, bins, points, features)
        # размерность тензора p повторяет первые размерности входного тензора x, например, (batch, height, width, bins * features)
        # в случае reduction='mean' размерность points сокращается, т.е. (batch, bins * features)
        return self._reduce(z_moved)
    def _forward_binary(self, x: torch.Tensor) -> torch.Tensor: #BinarySoftOrdering
        return self._reduce(self.sigmoid(x))
    def _reduction_pass(self, x: torch.Tensor) -> torch.Tensor:
        return x.flatten(-4, -3)
    def _reduction_mean(self, x: torch.Tensor) -> torch.Tensor:
        return x.mean(-3) 


class BinarySoftOrdering(nn.Module):  #Бинарный слой мягкого упорядочивания
    def __init__(self, reduction: str = 'mean') -> None:
        super().__init__()
        self.reduction = reduction
        self.sigmoid = nn.Sigmoid()
        if reduction == 'none':
            self.forward = self.forward_pass
        elif reduction == 'mean':
            self.forward = self.forward_mean
        else:
            assert "Unknown reduction function. Use 'none' or 'mean'."
    def extra_repr(self) -> str:
        if self.reduction == 'none':
            info = "reduction=None (Soft One-Hot Encoding)"
        else:
            info = "reduction=mean (Soft Ordering)"
        return info
    def forward_pass(self, x: torch.Tensor) -> torch.Tensor: # (batch, points, features)
        return self.sigmoid(x) # (batch, points, features)
    def forward_mean(self, x: torch.Tensor) -> torch.Tensor: # (batch, points, features)
        return self.sigmoid(x).mean(-2) # (batch, features)


# tabular-playground-series-apr-2022 [PyTorch] Bidirectional LSTM, ONN3
class PositionalSoftOrdering(nn.Module):
    """
    Слой, комбинирующий SOHE для признаков и позиций с независимыми
    гиперпараметрами, с последующим суммированием матриц взаимодействия.
    ОПТИМИЗИРОВАН для CUDA с использованием torch.einsum и torch.amp.
    """
    def __init__(self, 
                 n_bins_features: int,
                 n_bins_pos: int,
                 n_heads: int,
                 gcoef_features: float = 6.0,
                 gcoef_pos: float = 6.0,
                 use_amp: bool = True): # ОПТИМИЗАЦИЯ: Флаг для включения смешанной точности
        super().__init__()
        self.n_bins_features = n_bins_features
        self.gcoef_features = gcoef_features
        self.n_heads = n_heads        
        self.n_bins_pos = n_bins_pos
        self.gcoef_pos = gcoef_pos
        self.use_amp = use_amp # ОПТИМИЗАЦИЯ: Сохраняем флаг

        # Инициализируем n_features как неопределенный (ленивая реализация)
        self.n_features = None

        # Слой для кодирования признаков с n_bins_features
        self.feature_sohe = SoftOrdering(n_bins=n_bins_features, gcoef=gcoef_features, reduction='none', input_sigmoid=True)

        # Компоненты для позиционного кодирования
        self.pos_fc = None  # pos_fc "ленивый"
        # Слой SOHE для позиций с n_bins_pos
        self.pos_sohe = SoftOrdering(n_bins=n_bins_pos, gcoef=gcoef_pos, reduction='none', input_sigmoid=True)
    
    def extra_repr(self) -> str:
        # Отображаем информацию о слое, включая n_heads
        info = (f"n_features={'UNINITIALIZED' if self.n_features is None else self.n_features}, "
                f"n_heads={self.n_heads}, "
                f"n_bins_features={self.n_bins_features}, "
                f"n_bins_pos={self.n_bins_pos}")
        return info

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        :param x: Входной тензор. Форма (batch_size, seq_len, n_features).
        :return: Тензор после кодирования. Форма: (batch, n_heads * n_bins_pos * n_bins_features * n_features).
        """
        # ОПТИМИЗАЦИЯ: Используем контекст autocast для смешанной точности (float16)
        # Это вдвое сокращает память и ускоряет вычисления на совместимых GPU
        with torch.cuda.amp.autocast(enabled=self.use_amp and x.is_cuda):
            if x.dim() == 2:
                x = x.unsqueeze(-1)
            
            # --- Ленивая инициализация ---
            if self.n_features is None:
                self.n_features = x.shape[-1] # Определяем и сохраняем размерность признаков при первом проходе
                self.pos_fc = nn.Linear(self.n_features + 1, self.n_heads, device=x.device, dtype=x.dtype)
            else:
                # Проверяем, что размерность признаков не изменилась
                assert x.shape[-1] == self.n_features, \
                    f"Input n_features changed after initialization. Expected {self.n_features}, got {x.shape[-1]}."
            
            batch_size, seq_len, _ = x.shape
            N_data = self.n_bins_features * self.n_features

            # 1. Кодирование признаков
            x_sohe = self.feature_sohe(x) #(batch, seq_len, N_data)

            # 2. Позиционное кодирование
            positions = torch.arange(seq_len, device=x.device, dtype=x.dtype)
            positions_normalized = positions / (seq_len - 1 if seq_len > 1 else 1)
            positions_expanded = positions_normalized.view(1, -1, 1).expand(batch_size, -1, -1) #(batch_size, seq_len, 1)
            pos_fc_input = torch.cat([x, positions_expanded], dim=-1) # Конкатенируем данные (x) и позиции
            pos_signal = self.pos_fc(pos_fc_input) #(batch, seq_len, n_heads)
            pos_encoding_raw = self.pos_sohe(pos_signal) #(batch, seq_len, n_heads * n_bins_pos)
            
            # ОПТИМИЗАЦИЯ: Меняем форму, чтобы разделить головы и бины для позиций
            # Это необходимо для корректной работы einsum
            pos_encoding = pos_encoding_raw.view(batch_size, seq_len, self.n_heads, self.n_bins_pos)

            # 3. Объединение и агрегация с помощью einsum
            # ОПТИМИЗАЦИЯ: Заменяем создание гигантского тензора на одну эффективную операцию einsum.
            # 'bsh, bsd -> bhd' означает:
            # - Взять тензоры с индексами (b)atch, (s)equence, (h)eads, (d)ata
            # - Перемножить их по совпадающим индексам b и s
            # - Просуммировать по индексу s (sequence)
            # - На выходе получить тензор с индексами b, h, d
            # Мы используем 'bshp' для pos_encoding, где 'p' - это n_bins_pos
            aggregated = torch.einsum('bshp, bsd -> bhpd', pos_encoding, x_sohe)
            # Форма aggregated: (batch, n_heads, n_bins_pos, N_data)

            # 4. Выпрямление результата
            # flatten со start_dim=1 выпрямляет все размерности, кроме первой (batch)
            return aggregated.flatten(start_dim=1)

#tabular-playground-series-apr-2022 [PyTorch] Bidirectional LSTM, ONN2.py
class RecurrentSoftOrdering(nn.Module):
    """
    Оптимизированная версия RecurrentSoftOrdering с использованием JIT-компилятора.

    Изменения:
    1. `@torch.jit.script`: Компилирует весь модуль, включая циклы, в 
       оптимизированный граф для CUDA, устраняя узкие места Python.
    2. Удаление `.to(device)`: Модуль теперь автоматически работает на том
       устройстве, на котором находятся входные данные (`x.device`).
    3. Тип данных `float32`: Удалены преобразования в `.double()` для 
       стандартного и быстрого использования на GPU.
    4. Очистка кода: Удалены неиспользуемые переменные и закомментированные строки.
    5. Правильная инициализация тензоров: Тензоры создаются с `dtype` и `device`
       входного тензора `x`.
    """
    def __init__(self, n_inputs: int, n_hidden: int, n_bins: int, n_outputs: int, gcoef: float = 6.0, n_passes : int = 1, BiDirectional: bool = False) -> None:
        super().__init__()
        # JIT автоматически определяет константы, __constants__ не требуется
        self.n_inputs = n_inputs
        self.n_bins = n_bins
        self.n_outputs = n_outputs
        self.BiDirectional = BiDirectional

        self.linear1 = nn.Linear(n_bins * n_outputs + n_inputs, n_outputs)
        self.sigmoid1 = nn.Sigmoid()
        self.bn = nn.BatchNorm1d(n_outputs, affine=False, track_running_stats=False)  
        if n_bins >= 2:
            self.ordering = SoftOrdering(n_bins=n_bins, gcoef=gcoef, reduction='mean', input_sigmoid=True)
        else:
            self.ordering = BinarySoftOrdering(n_outputs, reduction='mean')

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Ожидаемый формат входа: (batch_size, seq_len, n_inputs)
        # Транспонируем для итерации по последовательности: (seq_len, batch_size, n_inputs)
        x = x.transpose(0, 1)
        
        seq_len, batch_size, _ = x.shape
        hidden_size = self.n_bins * self.n_outputs
        
        # --- Прямой проход ---
        # Инициализируем скрытое состояние на правильном устройстве и с правильным типом
        out1 = torch.zeros(batch_size, hidden_size, device=x.device, dtype=x.dtype)
        
        # JIT-компилятор оптимизирует этот цикл
        for i in range(seq_len):
            input_ = x[i]
            z = torch.cat((out1, input_), 1)
            z = self.linear1(z)
            #z = self.sigmoid1(z)
            z = z.unsqueeze(1)
            #z = self.bn(z)
            z = self.ordering(z)
            out1 = out1 + z / float(seq_len) # Используем float() для совместимости с JIT
        if self.BiDirectional:
            # --- Обратный проход ---
            out2 = torch.zeros(batch_size, hidden_size, device=x.device, dtype=x.dtype)
            
            # Итерация по перевернутой последовательности.
            # JIT также оптимизирует этот цикл.
            for i in range(seq_len):
                # Берем элементы с конца: x[seq_len - 1 - i]
                input_ = x[seq_len - 1 - i]
                z = torch.cat((out2, input_), 1)
                z = self.linear1(z)
                #z = self.sigmoid1(z)
                z = z.unsqueeze(1)
                z = self.bn(z)
                z = self.ordering(z)
                out2 = out2 + z / float(seq_len)

            # Объединяем результаты прямого и обратного проходов
            return torch.cat((out1, out2), 1)
        return out1
    def extra_repr(self) -> str:
        # n_hidden удален, так как не использовался
        return f"n_inputs={self.n_inputs}, n_bins={self.n_bins}, BiDirectional={self.BiDirectional}"


# To do: сделать по аналогии LazySoftOrdering - произвольную размерность данных, применить zslice для ускорения вычислений.
# MNIST_RaggingOrderingNN1.ipynb

class RaggingSoftOrdering(SoftOrdering):  #Сигмоидальный слой двухстороннего упорядочивания для тензоров переменной длины
    def forward(self, x: torch.Tensor, xlen: torch.Tensor) -> torch.Tensor:
        z = x.unsqueeze(2).repeat(1, 1, len(self.lspace), 1)
        l = self.lspace.unsqueeze(1).unsqueeze(0).unsqueeze(0).repeat(x.shape[0], x.shape[1], 1, x.shape[2])
        z = z - l
        z = self.sigmoid(self.gain * z)
        z = z[:, :, :-1] - z.roll(-1, dims=(2))[:, :, :-1]
        z = z.flatten()
        p = z.sum(1)
        p = p / xlen.unsqueeze(1).repeat(1, p.shape[1])
        return p

# BatchNorm чувствителен к нулевым хвостам в данных, так как "просматривает" весь батч,
# поэтому нужно сгенерировать усреднённые хвосты.
# Использование:
# z = average_tails(z, data_lengths)
# Нулевые хвосты нужны для линейных слоёв, сигмоидальных и др.:
# z = zero_tails(z, data_lengths)
# data_lengths - массив длин данных

# To do: хвосты нужно не просто усреднять, а генерировать с распределением,
# у которого не только среднее значение соответствует входным данным, но и разброс
# (стандартное отклонение)

def average_tails(data, data_lengths):
    pack_padded_sequence(data, data_lengths, batch_first=True, enforce_sorted=False)
    pad_sequence(data)
    m = torch.mean(data, (0, 1))
    pack_padded_sequence(data, data_lengths, batch_first=True, enforce_sorted=False)
    for k in range(data.shape[2]):
        pad_sequence(data[:, :, k], padding_value=m[k].item())
    return data

def zero_tails(data, data_lengths):
    pack_padded_sequence(data, data_lengths, batch_first=True, enforce_sorted=False)
    #pad_sequence(data)
    return data


# Применяется к входным данным, добавляет данные - обучаемые индексы
# Может переименовать в SoftPositionEncodingLayer???
# tabular-playground-series-apr-2022 [PyTorch] Bidirectional LSTM, ONN.ipynb
# tabular-playground-series-apr-2022 [PyTorch] Bidirectional LSTM, ONN2.ipynb

# To do: убрать batchnorm?

class IndexationLayer(nn.Module):  #Слой индексации
    len_Qvec: int
    n_vecs: int
    RelativeIndexes: bool
    TrainableStartIndexes: bool
    DynamicQuants: bool
    def __init__(self, len_Qvec: int, n_vecs: int, RelativeIndexes: bool = True, TrainableStartIndexes: bool = True, DynamicQuants: bool = True) -> None:
        super().__init__()
        __constants__ = ["len_Qvec", "n_vecs"]
        self.len_Qvec = len_Qvec
        self.n_vecs = n_vecs
        self.RelativeIndexes = RelativeIndexes
        self.TrainableStartIndexes = TrainableStartIndexes
        self.DynamicQuants = DynamicQuants
        self.sigm = nn.Sigmoid()
        self.bn = nn.LayerNorm(len_Qvec - 1)
        if RelativeIndexes:
            self.StartIndexes = torch.linspace(0, -1, steps=n_vecs).unsqueeze(1)
            if TrainableStartIndexes:
                self.StartIndexes = nn.Parameter(self.StartIndexes, requires_grad=True)
            else:
                self.StartIndexes.requires_grad = False
        else:
            self.StartIndexes = torch.zeros(self.n_vecs).unsqueeze(1)
        if DynamicQuants:
            self.Qvec = nn.Parameter(torch.randn(n_vecs, len_Qvec - 1, dtype=torch.double), requires_grad=True)
        else:
            self.Qvec = torch.zeros(n_vecs, len_Qvec - 1, dtype=torch.double)
            self.Qvec.requires_grad = False
            self.Q = torch.cumsum(self.sigm(self.Qvec), dim=1)
            self.Q = 0.5 * len_Qvec * F.normalize(self.Q, p=1, dim=1) + self.StartIndexes
            if not self.RelativeIndexes:
                self.RelativeQ = self.Q + self.StartIndexes 
    def extra_repr(self) -> str:
        return f"len_Qvec={self.len_Qvec}, n_vecs={self.n_vecs}, RelativeIndexes={self.RelativeIndexes}, TrainableStartIndexes={self.TrainableStartIndexes}, DynamicQuants={self.DynamicQuants}"
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.DynamicQuants:
            self.Q = torch.cumsum(self.sigm(self.Qvec), dim=1)
            ###self.Q = 0.5 * self.len_Qvec * F.normalize(self.Q, p=1, dim=1)
            #self.Q = self.bn(self.Q)
        #print(self.Qvec)
        #print(self.Q)
        if self.RelativeIndexes or self.DynamicQuants:
            self.RelativeQ = self.Q + self.StartIndexes 
        idx_tensor = torch.cat((self.StartIndexes, self.RelativeQ), 1).unsqueeze(0).unsqueeze(3).repeat(x.shape[0], 1, 1, 1)
        #print(idx_tensor.shape)
        #print(x.shape)
        x = x[:, :, 1:].unsqueeze(1).repeat(1, self.n_vecs, 1, 1)
        #print(idx_tensor.shape)
        x = torch.cat((x, idx_tensor), 3)
        return x