SpikingFlow.softbp.soft_pulse_function 源代码

import torch
import torch.nn as nn
import torch.nn.functional as F

class bilinear_leaky_relu(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, a=1, b=0.01, c=0.5):
        if x.requires_grad:
            piecewise0 = (x < -c).float()
            piecewise2 = (x > c).float()
            piecewise1 = torch.ones_like(x) - piecewise0 - piecewise2
            ctx.save_for_backward(piecewise0 * b + piecewise1 * a + piecewise2 * b)
        return (x >= 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = None
        if ctx.needs_input_grad[0]:
            grad_x = grad_output * ctx.saved_tensors[0]
        return grad_x, None, None, None

class BilinearLeakyReLU(nn.Module):
    def __init__(self, a=1, b=0.01, c=0.5):
        '''
        :param a: -c <= x <= c 时反向传播的梯度
        :param b: x > c 或 x < -c 时反向传播的梯度
        :param c: 决定梯度区间的参数
        :return: 与输入相同shape的输出

        双线性的近似脉冲发放函数。梯度为

        .. math::
            g'(x) =
            \\begin{cases}
            a, & -c \\leq x \\leq c \\\\
            b, & x < -c ~or~ x > c
            \\end{cases}

        对应的原函数为

        .. math::
            g(x) =
            \\begin{cases}
            bx + bc - ac, & x < -c \\\\
            ax, & -c \\leq x \\leq c \\\\
            bx - bc + ac, & x > c \\\\
            \\end{cases}

        '''
        super().__init__()
        self.a = a
        self.b = b
        self.c = c
        self.f = bilinear_leaky_relu.apply
    def forward(self, x):
        return self.f(x, self.a, self.b, self.c)

class sigmoid(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, alpha):
        if x.requires_grad:
            alpha_x = x * alpha
            ctx.save_for_backward(alpha_x)
            ctx.alpha = alpha
        return (x >= 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = None
        if ctx.needs_input_grad[0]:
            alpha_x = ctx.saved_tensors[0]
            s_x = torch.sigmoid(alpha_x)
            grad_x = grad_output * s_x * (1 - s_x) * ctx.alpha
        return grad_x, None

class Sigmoid(nn.Module):
    def __init__(self, alpha=1.0):
        '''
        :param x: 输入数据
        :param alpha: 控制反向传播时梯度的平滑程度的参数
        :return: 与输入相同shape的输出

        反向传播时使用sigmoid的梯度的脉冲发放函数。反向传播为

        .. math::
            g'(x) = \\alpha * (1 - \\mathrm{sigmoid} (\\alpha x)) \\mathrm{sigmoid} (\\alpha x)

        对应的原函数为

        .. math::
            g(x) = \\mathrm{sigmoid}(\\alpha x) = \\frac{1}{1+e^{-\\alpha x}}
        '''
        super().__init__()
        self.alpha = alpha
        self.f = sigmoid.apply
    def forward(self, x):
        return self.f(x, self.alpha)

class sign_swish(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, beta=1.0):
        if x.requires_grad:
            beta_x = beta * x
            ctx.save_for_backward(beta_x)
            ctx.beta = beta
        return (x >= 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        grad_x = None
        if ctx.needs_input_grad[0]:
            beta_x = ctx.saved_tensors[0]
            grad_x = ctx.beta * (2 - beta_x * torch.tanh(beta_x / 2)) / (1 + torch.cosh(beta_x)) \
                     * grad_output

        return grad_x, None

class SignSwish(nn.Module):
    def __init__(self, beta=5.0):
        '''
        :param x: 输入数据
        :param beta: 控制反向传播的参数
        :return: 与输入相同shape的输出

        Darabi, Sajad, et al. "BNN+: Improved binary network training." arXiv preprint arXiv:1812.11800 (2018).

        反向传播时使用swish的梯度的脉冲发放函数。反向传播为

        .. math::
            g'(x) = \\frac{\\beta (2 - \\beta x \\mathrm{tanh} \\frac{\\beta x}{2})}{1 + \\mathrm{cosh}(\\beta x)}

        对应的原函数为

        .. math::
            g(x) = 2 * \\mathrm{sigmoid}(\\beta x) * (1 + \\beta x (1 - \\mathrm{sigmoid}(\\beta x))) - 1
        '''
        super().__init__()
        self.beta = beta
        self.f = sign_swish.apply
    def forward(self, x):
        return self.f(x, self.beta)