dsipts 1.1.5__py3-none-any.whl

dsipts/models/duet/layers.py

@@ -0,0 +1,438 @@
+import torch
+import torch.nn as nn
+from torch.distributions.normal import Normal
+from einops import rearrange
+from ..autoformer.layers import series_decomp
+
+class encoder(nn.Module):
+    def __init__(self, seq_len,num_experts,hidden_size ):
+        super(encoder, self).__init__()
+        input_size = seq_len
+        num_experts = num_experts
+        encoder_hidden_size = hidden_size
+
+        self.distribution_fit = nn.Sequential(nn.Linear(input_size, encoder_hidden_size, bias=False), nn.ReLU(),
+                                              nn.Linear(encoder_hidden_size, num_experts, bias=False))
+
+    def forward(self, x):
+        mean = torch.mean(x, dim=-1)
+        out = self.distribution_fit(mean)
+        return out
+
+
+class RevIN(nn.Module):
+    def __init__(self, num_features: int, eps=1e-5, affine=True, subtract_last=False):
+        """
+        :param num_features: the number of features or channels
+        :param eps: a value added for numerical stability
+        :param affine: if True, RevIN has learnable affine parameters
+        """
+        super(RevIN, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.affine = affine
+        self.subtract_last = subtract_last
+        if self.affine:
+            self._init_params()
+
+    def forward(self, x, mode: str):
+        if mode == 'norm':
+            self._get_statistics(x)
+            x = self._normalize(x)
+        elif mode == 'denorm':
+            x = self._denormalize(x)
+        else:
+            raise NotImplementedError
+        return x
+
+    def _init_params(self):
+        # initialize RevIN params: (C,)
+        self.affine_weight = nn.Parameter(torch.ones(self.num_features))
+        self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
+
+    def _get_statistics(self, x):
+        dim2reduce = tuple(range(1, x.ndim - 1))
+        if self.subtract_last:
+            self.last = x[:, -1, :].unsqueeze(1)
+        else:
+            self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()  
+        self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
+
+    def _normalize(self, x):
+        if self.subtract_last:
+            x = x - self.last
+        else:
+            x = x - self.mean
+        x = x / self.stdev
+        if self.affine:
+            x = x * self.affine_weight
+            x = x + self.affine_bias
+        return x
+
+    def _denormalize(self, x):
+        if self.affine:
+            x = x - self.affine_bias
+            x = x / (self.affine_weight + self.eps * self.eps)
+        x = x * self.stdev
+        if self.subtract_last:
+            x = x + self.last
+        else:
+            x = x + self.mean
+        return x
+
+
+class Linear_extractor(nn.Module):
+    """
+    Paper link: https://arxiv.org/pdf/2205.13504.pdf
+    """
+
+    def __init__(self, seq_len,d_model,enc_in,CI,moving_avg, individual=False):
+        """
+        individual: Bool, whether shared model among different variates.
+        """
+        super(Linear_extractor, self).__init__()
+        self.seq_len = seq_len
+
+        self.pred_len = d_model
+        self.decompsition = series_decomp(moving_avg)
+        self.individual = individual
+        self.channels = enc_in
+        self.enc_in = 1 if CI else enc_in
+        if self.individual:
+            self.Linear_Seasonal = nn.ModuleList()
+            self.Linear_Trend = nn.ModuleList()
+
+            for i in range(self.channels):
+                self.Linear_Seasonal.append(
+                    nn.Linear(self.seq_len, self.pred_len))
+                self.Linear_Trend.append(
+                    nn.Linear(self.seq_len, self.pred_len))
+
+                self.Linear_Seasonal[i].weight = nn.Parameter(
+                    (1 / self.seq_len) * torch.ones([self.pred_len, self.seq_len]))
+                self.Linear_Trend[i].weight = nn.Parameter(
+                    (1 / self.seq_len) * torch.ones([self.pred_len, self.seq_len]))
+        else:
+            self.Linear_Seasonal = nn.Linear(self.seq_len, self.pred_len)
+            self.Linear_Trend = nn.Linear(self.seq_len, self.pred_len)
+
+            self.Linear_Seasonal.weight = nn.Parameter(
+                (1 / self.seq_len) * torch.ones([self.pred_len, self.seq_len]))
+            self.Linear_Trend.weight = nn.Parameter(
+                (1 / self.seq_len) * torch.ones([self.pred_len, self.seq_len]))
+
+
+
+    def encoder(self, x):
+        seasonal_init, trend_init = self.decompsition(x)
+        seasonal_init, trend_init = seasonal_init.permute(
+            0, 2, 1), trend_init.permute(0, 2, 1)
+        if self.individual:
+            seasonal_output = torch.zeros([seasonal_init.size(0), seasonal_init.size(1), self.pred_len],
+                                          dtype=seasonal_init.dtype).to(seasonal_init.device)
+            trend_output = torch.zeros([trend_init.size(0), trend_init.size(1), self.pred_len],
+                                       dtype=trend_init.dtype).to(trend_init.device)
+            for i in range(self.channels):
+                seasonal_output[:, i, :] = self.Linear_Seasonal[i](
+                    seasonal_init[:, i, :])
+                trend_output[:, i, :] = self.Linear_Trend[i](
+                    trend_init[:, i, :])
+        else:
+            seasonal_output = self.Linear_Seasonal(seasonal_init)
+            trend_output = self.Linear_Trend(trend_init)
+        x = seasonal_output + trend_output
+        return x.permute(0, 2, 1)
+
+    def forecast(self, x_enc):
+        # Encoder
+        return self.encoder(x_enc)
+
+
+    def forward(self, x_enc):
+        if x_enc.shape[0] == 0:
+            return torch.empty((0, self.pred_len, self.enc_in)).to(x_enc.device)
+        dec_out = self.forecast(x_enc)
+        return dec_out[:, -self.pred_len:, :]  # [B, L, D]
+
+
+class SparseDispatcher(object):
+    """Helper for implementing a mixture of experts.
+    The purpose of this class is to create input mini-batches for the
+    experts and to combine the results of the experts to form a unified
+    output tensor.
+    There are two functions:
+    dispatch - take an input Tensor and create input Tensors for each expert.
+    combine - take output Tensors from each expert and form a combined output
+      Tensor.  Outputs from different experts for the same batch element are
+      summed together, weighted by the provided "gates".
+    The class is initialized with a "gates" Tensor, which specifies which
+    batch elements go to which experts, and the weights to use when combining
+    the outputs.  Batch element b is sent to expert e iff gates[b, e] != 0.
+    The inputs and outputs are all two-dimensional [batch, depth].
+    Caller is responsible for collapsing additional dimensions prior to
+    calling this class and reshaping the output to the original shape.
+    See common_layers.reshape_like().
+    Example use:
+    gates: a float32 `Tensor` with shape `[batch_size, num_experts]`
+    inputs: a float32 `Tensor` with shape `[batch_size, input_size]`
+    experts: a list of length `num_experts` containing sub-networks.
+    dispatcher = SparseDispatcher(num_experts, gates)
+    expert_inputs = dispatcher.dispatch(inputs)
+    expert_outputs = [experts[i](expert_inputs[i]) for i in range(num_experts)]
+    outputs = dispatcher.combine(expert_outputs)
+    The preceding code sets the output for a particular example b to:
+    output[b] = Sum_i(gates[b, i] * experts[i](inputs[b]))
+    This class takes advantage of sparsity in the gate matrix by including in the
+    `Tensor`s for expert i only the batch elements for which `gates[b, i] > 0`.
+    """
+
+    def __init__(self, num_experts, gates):
+        """Create a SparseDispatcher."""
+
+        self._gates = gates
+        self._num_experts = num_experts
+        # sort experts
+        sorted_experts, index_sorted_experts = torch.nonzero(gates).sort(0)
+        # drop indices
+        _, self._expert_index = sorted_experts.split(1, dim=1)
+        # get according batch index for each expert
+        self._batch_index = torch.nonzero(gates)[index_sorted_experts[:, 1], 0]
+        # calculate num samples that each expert gets
+        self._part_sizes = (gates > 0).sum(0).tolist()
+        # expand gates to match with self._batch_index
+        gates_exp = gates[self._batch_index.flatten()]
+        self._nonzero_gates = torch.gather(gates_exp, 1, self._expert_index)
+
+    def dispatch(self, inp):
+        """Create one input Tensor for each expert.
+        The `Tensor` for an expert `i` contains the slices of `inp` corresponding
+        to the batch elements `b` where `gates[b, i] > 0`.
+        Args:
+          inp: a `Tensor` of shape "[batch_size, <extra_input_dims>]`
+        Returns:
+          a list of `num_experts` `Tensor`s with shapes
+            `[expert_batch_size_i, <extra_input_dims>]`.
+        """
+
+        # assigns samples to experts whose gate is nonzero
+
+        # expand according to batch index so we can just split by _part_sizes
+        inp_exp = inp[self._batch_index].squeeze(1)
+        return torch.split(inp_exp, self._part_sizes, dim=0)
+
+    def combine(self, expert_out, multiply_by_gates=True):
+        """Sum together the expert output, weighted by the gates.
+        The slice corresponding to a particular batch element `b` is computed
+        as the sum over all experts `i` of the expert output, weighted by the
+        corresponding gate values.  If `multiply_by_gates` is set to False, the
+        gate values are ignored.
+        Args:
+          expert_out: a list of `num_experts` `Tensor`s, each with shape
+            `[expert_batch_size_i, <extra_output_dims>]`.
+          multiply_by_gates: a boolean
+        Returns:
+          a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
+        """
+        # apply exp to expert outputs, so we are not longer in log space
+        stitched = torch.cat(expert_out, 0)
+        if multiply_by_gates:
+            # stitched = stitched.mul(self._nonzero_gates)
+            stitched = torch.einsum("i...,ij->i...", stitched, self._nonzero_gates)
+
+        shape = list(expert_out[-1].shape)
+        shape[0] = self._gates.size(0)
+        zeros = torch.zeros(*shape, requires_grad=True, device=stitched.device)
+        # combine samples that have been processed by the same k experts
+        combined = zeros.index_add(0, self._batch_index, stitched.float())
+        return combined
+
+    def expert_to_gates(self):
+        """Gate values corresponding to the examples in the per-expert `Tensor`s.
+        Returns:
+          a list of `num_experts` one-dimensional `Tensor`s with type `tf.float32`
+              and shapes `[expert_batch_size_i]`
+        """
+        # split nonzero gates for each expert
+        return torch.split(self._nonzero_gates, self._part_sizes, dim=0)
+
+
+class Linear_extractor_cluster(nn.Module):
+    """Call a Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
+    Args:
+    input_size: integer - size of the input
+    output_size: integer - size of the input
+    num_experts: an integer - number of experts
+    hidden_size: an integer - hidden size of the experts
+    noisy_gating: a boolean
+    k: an integer - how many experts to use for each batch element
+    """
+
+    def __init__(self, noisy_gating,num_experts,seq_len,k,d_model,enc_in,CI,moving_avg,hidden_size ):
+        super(Linear_extractor_cluster, self).__init__()
+        self.noisy_gating = noisy_gating
+        self.num_experts = num_experts
+        self.input_size = seq_len
+        self.k = k
+        # instantiate experts
+        self.experts = nn.ModuleList([Linear_extractor(seq_len,d_model,enc_in,CI,moving_avg) for _ in range(self.num_experts)])
+        self.W_h = nn.Parameter(torch.eye(self.num_experts))
+        self.gate = encoder(seq_len,num_experts,hidden_size)
+        self.noise = encoder(seq_len,num_experts,hidden_size)
+
+        self.n_vars = enc_in
+        self.revin = RevIN(self.n_vars)
+
+        self.CI = CI
+        self.softplus = nn.Softplus()
+        self.softmax = nn.Softmax(1)
+        self.register_buffer("mean", torch.tensor([0.0]))
+        self.register_buffer("std", torch.tensor([1.0]))
+        assert self.k <= self.num_experts
+
+    def cv_squared(self, x):
+        """The squared coefficient of variation of a sample.
+        Useful as a loss to encourage a positive distribution to be more uniform.
+        Epsilons added for numerical stability.
+        Returns 0 for an empty Tensor.
+        Args:
+        x: a `Tensor`.
+        Returns:
+        a `Scalar`.
+        """
+        eps = 1e-10
+        # if only num_experts = 1
+
+        if x.shape[0] == 1:
+            return torch.tensor([0], device=x.device, dtype=x.dtype)
+        return x.float().var() / (x.float().mean() ** 2 + eps)
+
+    def _gates_to_load(self, gates):
+        """Compute the true load per expert, given the gates.
+        The load is the number of examples for which the corresponding gate is >0.
+        Args:
+        gates: a `Tensor` of shape [batch_size, n]
+        Returns:
+        a float32 `Tensor` of shape [n]
+        """
+        return (gates > 0).sum(0)
+
+    def _prob_in_top_k(
+        self, clean_values, noisy_values, noise_stddev, noisy_top_values
+    ):
+        """Helper function to NoisyTopKGating.
+        Computes the probability that value is in top k, given different random noise.
+        This gives us a way of backpropagating from a loss that balances the number
+        of times each expert is in the top k experts per example.
+        In the case of no noise, pass in None for noise_stddev, and the result will
+        not be differentiable.
+        Args:
+        clean_values: a `Tensor` of shape [batch, n].
+        noisy_values: a `Tensor` of shape [batch, n].  Equal to clean values plus
+          normally distributed noise with standard deviation noise_stddev.
+        noise_stddev: a `Tensor` of shape [batch, n], or None
+        noisy_top_values: a `Tensor` of shape [batch, m].
+           "values" Output of tf.top_k(noisy_top_values, m).  m >= k+1
+        Returns:
+        a `Tensor` of shape [batch, n].
+        """
+        batch = clean_values.size(0)
+        m = noisy_top_values.size(1)
+        top_values_flat = noisy_top_values.flatten()
+
+        threshold_positions_if_in = (
+            torch.arange(batch, device=clean_values.device) * m + self.k
+        )
+        threshold_if_in = torch.unsqueeze(
+            torch.gather(top_values_flat, 0, threshold_positions_if_in), 1
+        )
+        is_in = torch.gt(noisy_values, threshold_if_in)
+        threshold_positions_if_out = threshold_positions_if_in - 1
+        threshold_if_out = torch.unsqueeze(
+            torch.gather(top_values_flat, 0, threshold_positions_if_out), 1
+        )
+        # is each value currently in the top k.
+        normal = Normal(self.mean, self.std)
+        prob_if_in = normal.cdf((clean_values - threshold_if_in) / noise_stddev)
+        prob_if_out = normal.cdf((clean_values - threshold_if_out) / noise_stddev)
+        prob = torch.where(is_in, prob_if_in, prob_if_out)
+        return prob
+
+    def noisy_top_k_gating(self, x, train, noise_epsilon=1e-2):
+        """Noisy top-k gating.
+        See paper: https://arxiv.org/abs/1701.06538.
+        Args:
+          x: input Tensor with shape [batch_size, input_size]
+          train: a boolean - we only add noise at training time.
+          noise_epsilon: a float
+        Returns:
+          gates: a Tensor with shape [batch_size, num_experts]
+          load: a Tensor with shape [num_experts]
+        """
+        clean_logits = self.gate(x)
+
+        if self.noisy_gating and train:
+            raw_noise_stddev = self.noise(x)
+            noise_stddev = self.softplus(raw_noise_stddev) + noise_epsilon
+            noise = torch.randn_like(clean_logits)
+            noisy_logits = clean_logits + (noise * noise_stddev)
+            logits = noisy_logits @ self.W_h
+        else:
+            logits = clean_logits
+
+        # calculate topk + 1 that will be needed for the noisy gates
+        logits = self.softmax(logits)
+        top_logits, top_indices = logits.topk(min(self.k + 1, self.num_experts), dim=1)
+        top_k_logits = top_logits[:, : self.k]
+        top_k_indices = top_indices[:, : self.k]
+        top_k_gates = top_k_logits / (
+            top_k_logits.sum(1, keepdim=True) + 1e-6
+        )  # normalization
+
+        zeros = torch.zeros_like(logits, requires_grad=True)
+        gates = zeros.scatter(1, top_k_indices, top_k_gates)
+
+        if self.noisy_gating and self.k < self.num_experts and train:
+            load = (
+                self._prob_in_top_k(
+                    clean_logits, noisy_logits, noise_stddev, top_logits
+                )
+            ).sum(0)
+        else:
+            load = self._gates_to_load(gates)
+        return gates, load
+
+    def forward(self, x, loss_coef=1):
+        """Args:
+        x: tensor shape [batch_size, input_size]
+        train: a boolean scalar.
+        loss_coef: a scalar - multiplier on load-balancing losses
+
+        Returns:
+        y: a tensor with shape [batch_size, output_size].
+        extra_training_loss: a scalar.  This should be added into the overall
+        training loss of the model.  The backpropagation of this loss
+        encourages all experts to be approximately equally used across a batch.
+        """
+        gates, load = self.noisy_top_k_gating(x, self.training)
+        # calculate importance loss
+        importance = gates.sum(0)
+        loss = self.cv_squared(importance) + self.cv_squared(load)
+        loss *= loss_coef
+
+        dispatcher = SparseDispatcher(self.num_experts, gates)
+        if self.CI:
+            x_norm = rearrange(x, "(x y) l c -> x l (y c)", y=self.n_vars)
+            x_norm = self.revin(x_norm, "norm")
+            x_norm = rearrange(x_norm, "x l (y c) -> (x y) l c", y=self.n_vars)
+        else:
+            x_norm = self.revin(x, "norm")
+
+        expert_inputs = dispatcher.dispatch(x_norm)
+
+        gates = dispatcher.expert_to_gates()
+        expert_outputs = [
+            self.experts[i](expert_inputs[i]) for i in range(self.num_experts)
+        ]
+        y = dispatcher.combine(expert_outputs)
+
+        return y, loss

dsipts/models/duet/masked.py

@@ -0,0 +1,202 @@
+import torch.nn as nn
+import torch
+from math import sqrt
+
+import torch.nn.functional as F
+from torch.nn.functional import gumbel_softmax
+import math
+import torch.fft
+from einops import rearrange
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation="relu"):
+        super(EncoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.attention = attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, attn_mask=None, tau=None, delta=None):
+        new_x, attn = self.attention(
+            x, x, x,
+            attn_mask=attn_mask,
+            tau=tau, delta=delta
+        )
+        x = x + self.dropout(new_x)
+
+        y = x = self.norm1(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm2(x + y), attn
+
+
+class Encoder(nn.Module):
+    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+        super(Encoder, self).__init__()
+        self.attn_layers = nn.ModuleList(attn_layers)
+        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+        self.norm = norm_layer
+
+    def forward(self, x, attn_mask=None, tau=None, delta=None):
+        # x [B, L, D]
+        attns = []
+        if self.conv_layers is not None:
+            for i, (attn_layer, conv_layer) in enumerate(zip(self.attn_layers, self.conv_layers)):
+                delta = delta if i == 0 else None
+                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)
+                x = conv_layer(x)
+                attns.append(attn)
+            x, attn = self.attn_layers[-1](x, tau=tau, delta=None)
+            attns.append(attn)
+        else:
+            for attn_layer in self.attn_layers:
+                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)
+                attns.append(attn)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x, attns
+
+
+class FullAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
+        super(FullAttention, self).__init__()
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        scale = self.scale or 1. / sqrt(E)
+
+        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+
+        # if self.mask_flag:
+        #     large_negative = -math.log(1e10)
+        #     attention_mask = torch.where(attn_mask == 0, torch.tensor(large_negative), attn_mask)
+        #
+        #     scores = scores * attention_mask
+        if self.mask_flag:
+            large_negative = -math.log(1e10)
+            attention_mask = torch.where(attn_mask == 0, large_negative, 0)
+
+            scores = scores * attn_mask + attention_mask
+
+        A = self.dropout(torch.softmax(scale * scores, dim=-1))
+        V = torch.einsum("bhls,bshd->blhd", A, values)
+
+        if self.output_attention:
+            return V.contiguous(), A
+        else:
+            return V.contiguous(), None
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, attention, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AttentionLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_attention = attention
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out, attn = self.inner_attention(
+            queries,
+            keys,
+            values,
+            attn_mask,
+            tau=tau,
+            delta=delta
+        )
+        out = out.view(B, L, -1)
+
+        return self.out_projection(out), attn
+
+
+class Mahalanobis_mask(nn.Module):
+    def __init__(self, input_size):
+        super(Mahalanobis_mask, self).__init__()
+        frequency_size = input_size // 2 + 1
+        self.A = nn.Parameter(torch.randn(frequency_size, frequency_size), requires_grad=True)
+
+    def calculate_prob_distance(self, X):
+        XF = torch.abs(torch.fft.rfft(X, dim=-1))
+        X1 = XF.unsqueeze(2)
+        X2 = XF.unsqueeze(1)
+
+        # B x C x C x D
+        diff = X1 - X2
+
+        temp = torch.einsum("dk,bxck->bxcd", self.A, diff)
+
+        dist = torch.einsum("bxcd,bxcd->bxc", temp, temp)
+
+        # exp_dist = torch.exp(-dist)
+        exp_dist = 1 / (dist + 1e-10)
+        # 对角线置零
+
+        identity_matrices = 1 - torch.eye(exp_dist.shape[-1])
+        mask = identity_matrices.repeat(exp_dist.shape[0], 1, 1).to(exp_dist.device)
+        exp_dist = torch.einsum("bxc,bxc->bxc", exp_dist, mask)
+        exp_max, _ = torch.max(exp_dist, dim=-1, keepdim=True)
+        exp_max = exp_max.detach()
+
+        # B x C x C
+        p = exp_dist / exp_max
+
+        identity_matrices = torch.eye(p.shape[-1])
+        p1 = torch.einsum("bxc,bxc->bxc", p, mask)
+
+        diag = identity_matrices.repeat(p.shape[0], 1, 1).to(p.device)
+        p = (p1 + diag) * 0.99
+
+        return p
+
+    def bernoulli_gumbel_rsample(self, distribution_matrix):
+        b, c, d = distribution_matrix.shape
+
+        flatten_matrix = rearrange(distribution_matrix, 'b c d -> (b c d) 1')
+        r_flatten_matrix = 1 - flatten_matrix
+
+        log_flatten_matrix = torch.log(flatten_matrix / r_flatten_matrix)
+        log_r_flatten_matrix = torch.log(r_flatten_matrix / flatten_matrix)
+
+        new_matrix = torch.concat([log_flatten_matrix, log_r_flatten_matrix], dim=-1)
+        resample_matrix = gumbel_softmax(new_matrix, hard=True)
+
+        resample_matrix = rearrange(resample_matrix[..., 0], '(b c d) -> b c d', b=b, c=c, d=d)
+        return resample_matrix
+
+    def forward(self, X):
+        p = self.calculate_prob_distance(X)
+
+        # bernoulli中两个通道有关系的概率
+        sample = self.bernoulli_gumbel_rsample(p)
+
+        mask = sample.unsqueeze(1)
+        cnt = torch.sum(mask, dim=-1)
+        return mask