胶囊网络中挤压函数Squash的实现（tensorflow和pytorch）

写在前面

最近参加学校的一个比赛，与组内成员讨论了一番准备使用胶囊网络实现对街景交通标志的识别，我们找到的基本的源代码是tensorflow实现的，需要将其转换成pytorch的代码（工作量巨大，草【一种植物】）
本文就胶囊网络的挤压函数的转化进行说明，记录，一方面是为了让自己记住，另一方面也能为后来的同学们提供方法。
另外这篇文章也会记录一些Latex的输入公式，主要是用于我自己看的哈哈哈

向量的二范数

$\Vert  V \Vert_2$ = $\sqrt{x_1^2+x_2^2+...+x_n^2}$，其中$x_i\in V$

$\Vert V{\Vert }_2$ = $\sqrt{x_1^2+x_2^2+...+x_n^2}$，其中$x_i\in V$
那么注意到其实向量二范数就是向量的模长，由各个元素的平方和开根号得到

挤压函数

$||\vec V_j^l||$=squash($||S_j^i||$)=$\frac{||S^l_j||^2}{1+||S^l_j||^2}\cdot\frac{S_j^l}{||S^l_j||^2}$

在文章中我们假设得到的向量是$||S||$,经过挤压后得到向量$||V||$
那么有以下公式：
$||{\vec{V}}_j^l||$=squash($||S_j^l||$)=$\frac{||S_j^l|{|}_2^2}{1+||S_j^l|{|}_2^2}\cdot \frac{S_j^l}{||S_j^l|{|}_2}$

挤压函数的实现

通过上面的公式我们知道，要想实现Squash函数，需要做平方和运算，开方运算。
tensorflow实现方式：

def squash(vector):#vector对应公式中的S
	'''
		tf.square:对输入变量的每个元素求平方
		tf.reduce_sum:第一个参数为待求和的向量，第二个参数为维度，第三个参数为是否改变输入向量的维度
	'''
	vector +=0.00001
	norm = tf.reduce_sum(tf.square(vector), -2, keep_dims=True)
	#这里的norm对应公式中S二范数的平方
	scalar = norm/(1+norm)/tf.sqrt(norm)
	V = scalar * vector#V对应公式中的V
	return V

pytorch实现方式

def squash(inputs, axis=-1):
	#这里的norm对应公式中S的二范数（跟上面的norm不一样）
	#torch.norm可以直接计算输入向量的范数，第二个参数为范数的类型，p=2为二范数
    norm = torch.norm(inputs, p=2, dim=axis, keepdim=True)
    scale = norm**2 / (1 + norm**2) /norm 
    V = scale * inputs
    return V