


    def __init__(self):
        super(SplitValuePooling, self).__init__()
        # self.gap = nn.AdaptiveAvgPool2d(1)
        self.gap1 = GeneralizedMeanPoolingP()
        self.gap2 = GeneralizedMeanPoolingP()

    def forward(self, x):
        # < 0
        x1 = torch.max(-x, other=torch.zeros_like(x, device=x.device))
        x1 = self.gap1(x1)
        # > 0
        x2 = torch.max(x, other=torch.zeros_like(x, device=x.device))
        x2 = self.gap2(x2)

        return torch.cat([x1, x2], dim=1)


class NonLocal(nn.Module):
    def __init__(self, in_channels, inter_channels=None, bn_layer=False):
        super(NonLocal, self).__init__()


        self.in_channels = in_channels
        self.inter_channels = inter_channels

        if self.inter_channels is None:
            self.inter_channels = in_channels // 4
            if self.inter_channels == 0:
                self.inter_channels = 1


        self.g = nn.Conv2d(in_channels=self.in_channels,
                         out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0)

        if bn_layer:
            self.W = nn.Sequential(
                nn.Conv2d(in_channels=self.inter_channels,
                          out_channels=self.in_channels,
                        kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(self.in_channels)
            )
            nn.init.constant_(self.W[1].weight, 0)
            nn.init.constant_(self.W[1].bias, 0)
        else:
            self.W = nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels,
                             kernel_size=1, stride=1, padding=0)
            nn.init.constant_(self.W.weight, 0)
            # nn.init.constant_(self.W.bias, 0)

        self.theta = nn.Conv2d(in_channels=self.in_channels,
                               out_channels=self.inter_channels,
                              kernel_size=1, stride=1, padding=0)
        self.phi = nn.Conv2d(in_channels=self.in_channels,
                              out_channels=self.inter_channels,
                              kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        '''
        :param x: (b, c, t, h, w)
        :return:
        '''

        batch_size = x.size(0)

        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        f = torch.matmul(theta_x, phi_x)
        f_div_C = F.softmax(f, dim=-1)

        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        return z


class PairGraph(nn.Module):
    def __init__(self, in_channels, inter_channels=None):
        super(PairGraph, self).__init__()


        self.in_channels = in_channels
        self.inter_channels = inter_channels

        if self.inter_channels is None:
            self.inter_channels = in_channels // 16
            if self.inter_channels == 0:
                self.inter_channels = 1


        self.y = nn.Conv2d(in_channels=self.in_channels,
                         out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0, bias=False)

        self.g = nn.Conv2d(in_channels=self.in_channels,
                               out_channels=self.inter_channels,
                              kernel_size=1, stride=1, padding=0, bias=False)
        self.o = nn.Conv2d(in_channels=self.inter_channels,
                              out_channels=self.in_channels,
                              kernel_size=1, stride=1, padding=0, bias=False)

        nn.init.constant_(self.o.weight, 0)
        # nn.init.constant_(self.o.bias, 0)

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        '''
        :param x: (b, c, t, h, w)
        :return:
        '''

        b, _, h, w = x.size()

        y = self.y(x)
        y = normalize(y, axis=1)   # n, c, h, w
        y = y.reshape(b, self.inter_channels, -1)  # n, c, h*w
        y1 = y.permute(0, 2, 1)   # n, h*w, c

        coef = self.relu(y1.matmul(y))    # n, h*w, h*w

        batch_eye = torch.eye(h*w, device=x.device).reshape(1, h*w, h*w).repeat(b, 1, 1)
        adja = coef * (1. - batch_eye)
        degree = torch.sum(adja, dim=-1, keepdim=True) * batch_eye
        # flag = degree > 1e-3
        temp = torch.where(degree > 1e-3, degree, torch.ones_like(degree, device=x.device))
        # temp = flag * degree + (1 - flag) * torch.ones_like(degree, device=x.device)
        norm = temp.rsqrt() * batch_eye
        norm = norm.detach()
        f = torch.matmul(torch.matmul(norm, degree - adja), norm)   # n, h*w, h*w

        g = self.g(x).reshape(b, self.inter_channels, -1)  # n, c, h*w
        o = torch.matmul(g, f).reshape(b, self.inter_channels, h, w)
        o = self.o(o)
        return x + o


class DSBN2d(nn.Module):
    def __init__(self, planes):
        super(DSBN2d, self).__init__()
        self.num_features = planes
        self.BN_S = nn.BatchNorm2d(planes)
        self.BN_T = nn.BatchNorm2d(planes)

    def forward(self, x):
        if (not self.training):
            return self.BN_S(x)

        bs = x.size(0)
        assert (bs % 2 == 0)
        split = torch.split(x, int(bs/2), 0)
        out1 = self.BN_S(split[0].contiguous())
        out2 = self.BN_T(split[1].contiguous())
        out = torch.cat((out1, out2), 0)
        return out


class DSBN2dConstBatch(nn.Module):
    def __init__(self, planes, batch_size=64, constant_batch=32):
        super(DSBN2dConstBatch, self).__init__()
        self.num_features = planes
        self.constant_batch = constant_batch
        self.bn_list = nn.Sequential(*[nn.BatchNorm2d(planes) for _ in range(batch_size // constant_batch)])

