复数神经网络及其 PyTorch 实现

作者丨科技猛兽

编辑丨极市平台

极市导读

 

实数网络在图像领域取得极大成功，但在音频中，信号特征大多数是复数，如频谱等。简单分离实部虚部，或者考虑幅度和相位角都丢失了复数原本的关系。论文按照复数计算的定义，设计了深度复数网络，能对复数的输入数据进行卷积、激活、批规范化等操作。这里对论文提出的几种复数操作进行介绍，并给出简单的 Pytorch 实现方法。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

虽然叫深度复数网络，但里面的操作实际上还是在实数空间进行的。但通过实数的层实现类似于复数计算的操作。

目录

1 PyTorch 中的复数张量形式

2 复数神经网络背景

3 复数卷积操作
3.1 复数卷积原理
3.2 复数卷积 PyTorch 实现
3.3 复数的反向传播
3.4 柯西-黎曼方程 (Cauchy–Riemann Equation)

4 复数 LSTM 操作
4.1 复数 LSTM 原理
4.2 复数 LSTM PyTorch 实现

5 复数激活函数
5.1 复数激活函数原理
5.2 复数激活函数 PyTorch 实现

6 复数 Dropout
6.1 复数 Dropout原理
6.2 复数 Dropout PyTorch 实现

7 复数权重初始化
7.1 复数权重初始化原理

8 复数 Batch Normalization
8.1 复数 BN 原理
8.2 复数 BN PyTorch 实现

9 完整模型搭建

1 PyTorch 中的复数张量形式

PyTorch 1.8 及之后都支持2种复数形式的 Tensor，它们分别是：

意味着 torch 中有表示 complex 的张量形式，即：

torch.complex(real, imag, *, out=None) → Tensor

构造一个复数张量，其实部等于 real，虚部等于 imag。

Parameters

• real (Tensor)： 复数张量的实数部分。必须为 float 或 double。
• imag (Tensor)： 复数张量的虚部。dtype 必须与实部 real 相同。

关键字参数：

out (Tensor)： 如果输入为 torch.float32 ，则必须为 torch.complex64 。如果输入为 torch.float64 ，则必须为 torch.complex128 。

torch.is_complex(input)

返回 input 是不是复数形式，也就是torch.complex64, 和torch.complex128中的一种。

2 复数神经网络背景

众所周知， 从计算、生物和信号处理的角度来看，使用复数有许多优点。所以，复数相对于实数具有更强的表达能力。若能够借助复数设计神经网络，则非常具有吸引力。但是一个难题是如何设计配套的各种网络的 building block，比如说 complex BN，complex weight initialization 等等。

复数神经网络也有一些生物学上的优势，即：若网络中的数据都是实数，则只能代表某个中间输出的具体的值的大小；反之，若网络中的数据都是复数，则不仅能代表某个中间输出的具体的值的大小 (复数的模长)，还可以代表时间的概念 (复数的相位)。具有相似相位的输入神经元是同步的 (synchronous)，因为它们在复数运算中是相加的，而异步神经元相加则具有破坏性 (asynchronous)，因此相互干扰。

复数神经网络也有一些信号处理方面的优势，即：复数蕴含着相位信息，而语音信号中的相位信息影响其可懂度。奥本海姆的研究表明，在图像的相位中存在的信息量足以恢复以其幅值编码的大部分信息。事实上，相位信息在对物体的形状，边缘和方向进行编码时，提供了对物体的详细描述。

本文开发了适当的工具和一个通用的框架来训练具有复杂参数的深层神经网络。

3 复数卷积操作

3.1 复数卷积原理

任意的一个复数  ，其实部为  ，虚部为  。作者将复数的实部和虚部表示为逻辑上不同的实值实体，并在内部使用实值算术模拟复数运算。假设一个卷积核，权重是  ，则它可以表示成  个复数权重。

复数域上执行传统的实值二维卷积：

复数卷积核： 

复数输入张量： 

复数卷积过程： 

在具体实现中，可以使用下图1所示的简单结构实现。

图1：复数域上执行传统的实值二维卷积的过程

如下图1所示，把上式写成矩阵的形式，就有：

3.2 复数卷积 PyTorch 实现

PyTorch 实现复数的操作基于 apply_complex 这个方法。

def apply_complex(fr, fi, input, dtype = torch.complex64):
return (fr(input.real)-fi(input.imag)).type(dtype) \
+ 1j*(fr(input.imag)+fi(input.real)).type(dtype)


这个函数需要传入2个操作 (nn.Conv2d, nn.Linear 等等) 和 torch.complex64 类型的 input。
fr(input.real)： 卷积核的实部 * (输入的实部)。
fi(input.imag)： 卷积核的虚部 * (输入的虚部)
fr(input.imag)： 卷积核的实部 * (输入的虚部)
fi(input.real)： 卷积核的虚部 * (输入的实部)
input 类型： torch.complex64
返回值类型： torch.complex64

因此，利用 Pytorch 的 nn.Conv2D 实现，严格遵守上面复数卷积的定义式：

class ComplexConv2d(Module):

    def __init__(self,in_channels, out_channels, kernel_size=3, stride=1, padding = 0,
                 dilation=1, groups=1, bias=True):
        super(ComplexConv2d, self).__init__()
        self.conv_r = Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.conv_i = Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        
    def forward(self,input):    
        return apply_complex(self.conv_r, self.conv_i, input)

同理还可以实现 Pytorch 的 nn.Linear和 Pytorch 的 nn.ConvTranspose2d：

class ComplexLinear(Module):

def __init__(self, in_features, out_features):
super(ComplexLinear, self).__init__()
self.fc_r = Linear(in_features, out_features)
self.fc_i = Linear(in_features, out_features)

def forward(self, input):
return apply_complex(self.fc_r, self.fc_i, input)

class ComplexConvTranspose2d(Module):

def __init__(self,in_channels, out_channels, kernel_size, stride=1, padding=0,
output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros'):

super(ComplexConvTranspose2d, self).__init__()

self.conv_tran_r = ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding,
output_padding, groups, bias, dilation, padding_mode)
self.conv_tran_i = ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding,
output_padding, groups, bias, dilation, padding_mode)


具体实现的思路相似，都是借助了 apply_complex 函数，传入2个操作 (nn.Conv2d, nn.Linear 等等) 和 torch.complex64 类型的 input，然后在 ComplexLinear (或 ComplexConvTranspose2d) 中分别计算。

3.3 复数的反向传播

为了在复数神经网络中进行反向传播，一个充分条件是网络训练的目标函数和激活函数对网络中每个 complex parameter 的实部和虚部都是可微的。通常损失函数都是实数，则复数 chain rule 如下：

如果  是实数损失函数，  为复变量，满足  ，则有：

如果现在有另一个复数  ，且  ，则根据偏导数的链式法则：

3.4 柯西-黎曼方程 (Cauchy–Riemann Equation)

一个解析函数的解析性 (Holomorphism, analyticity) 确保一个复变函数在其域中的每个点的邻域中都是可微的。假设有复变函数  ，和一个复数  ，并有：  (和  是  实函数)。这意味着函数  的导数，即：  在  的邻域内都存在。

，也就是说  可以沿着不同的方向逼近0 (可以沿着实轴，虚轴或中间)。所以说当  沿着实轴方向逼近0时，有：

当  沿着虚轴方向逼近0时，有：

根据式4和式5我们可以推导出：

所以，为了使得上式恒成立，就必须满足下式：

式7就被称为柯西-黎曼方程 (Cauchy–Riemann Equation)。所以：

定理1：设函数  定义在区域  内，则  在  内一点  可导的充要条件是  与  在点  可微，并且在该点满足柯西-黎曼方程 (Cauchy–Riemann Equation)。

4 复数 LSTM 操作

卷积 LSTM 类似于完全连接 LSTM。唯一的区别是，没有使用矩阵乘法来执行计算，而是使用卷积运算。实值卷积 LSTM 的计算定义如下：

式中，  代表 sigmoid 激活函数，  代表 element-wise 的乘法，  代表实数卷积。  分别是 input gate，forget gate，和 output gate。  和  代表 cell 和 hidden states。

对于复卷积 LSTM，只需用复卷积运算代替实值卷积运算。elementwise multiplication 保持不变，Sigmoid 和 tanh 分别在实部和虚部进行。

5 复数激活函数

5.1 复数激活函数原理

复数激活函数在之前的工作里面已有研究：modReLU 和 zReLU。modReLU 可以表示为：

式中，  ，  是  的相位，  是可学习的参数。设置参数  的原因是  总是正的，所以参数  使得激活函数可以到达 dead zone 的位置。modReLU 的特点是激活函数前后复数的相位是不变的，但是modReLU 激活函数不满足柯西-黎曼方程 (Cauchy–Riemann Equations)，因此它不是解析的。而复数激活函数需要满足 Cauchy-Riemann Equations 才是解析的，才能进行复数微分操作。

modReLU 不满足 Cauchy-Riemann Equations。

接下来是 zReLU 激活函数。

zReLU 在  的时候不满足，即在x和y的正半轴不满足 Cauchy-Riemann Equations。

接下来是 CReLU 激活函数。它的设计初衷是要满足柯西-黎曼方程 (Cauchy–Riemann Equation)，所以 CReLU 激活函数分别在实部和虚部上进行激活操作，即：

当实部和虚部同时为严格正或严格负时，CReLU 激活函数满足 Cauchy-Riemann Equations，这里我给个简单证明。

证明： 设  ，则  。

得证。

CReLU 只在实部虚部同时大于零或同时小于零的时候满足 Cauchy-Riemann Equations，即在第2,4象限不满足。

5.2 复数激活函数 PyTorch 实现

from torch.nn.functional import relu

def complex_relu(input):
    return relu(input.real).type(torch.complex64)+1j*relu(input.imag).type(torch.complex64)

6 复数 Dropout

6.1 复数 Dropout 原理

复数 Dropout 个人感觉实部虚部需要同时置0，作者源码中没用到 Dropout 层。

6.2 复数 Dropout PyTorch 实现

from torch.nn.functional import dropout

def complex_dropout(input, p=0.5, training=True):
# need to have the same dropout mask for real and imaginary part,
mask = torch.ones(*input.shape, dtype = torch.float32)
mask = dropout(mask, p, training)*1/(1-p)
mask.type(input.dtype)
return mask*input


7 复数权重初始化

7.1 复数权重初始化原理

作者介绍了两种初始化方法的复数形式：Glorot，Kaiming 初始化，复数形式的权重可以表示为：

式中，  和  分别是权重的幅值和相位。权重的方差可以这样计算：

根据以上2式有：

当一个随机二维向量的两个分量呈独立的、均值为0，有着相同的方差的正态分布时，这个向量的模呈瑞利分布 (Rayleigh Distribution)。

瑞利分布简介

设  和  是相互独立的随机变量，并且均服从零均值的高斯分布：

有一个新的随机变量  ，它与  和  是的关系是： ，则变量  服从瑞利分布。

我们推导一下变量  的概率分布：

代表  和  的联合分布在以  为半径的圆内的概率。而  和  的分布是独立的高斯分布，则有：

所以变量  的概率分布是  。

所以复数权重的模型  的分布服从期望为  ，方差为  的瑞利分布，即参数为  的瑞利分布。我们进一步有：

。

如果我们按照 Glorot 初始化的方式，则需要满足  ，此时有：  。如果我们按照 Kaiming 初始化的方式，则需要满足  ，此时有：  。

到这里发现  的大小与  的幅值有关，而与具体的相位无关。所以利用  之间的均匀分布来初始化相位。通过执行式12中的相量乘以幅值，就完成了复数权重的初始化。

8 复数 Batch Normalization

8.1 复数 BN 原理

首先肯定不能用常规的 BN 方法，否则实部和虚部的分布就不能保证了。但正如常规 BN 方法，首先要对输入进行0均值1方差的操作，只是方法有所不同。复数 BN 的方法是：

协方差矩阵  是：

通过15式的操作，可以确保输出的均值为0，协方差为1，相关为0。这里有一点需要特别注意的是求协方差矩阵  的逆矩阵。

输入要乘以  的逆平方根，即：

BN 中还有  和  两个参数，因此最终结果如下：

因为归一化的输入  的模长是1，所以  和  被初始化为  。而  ,  ,  ,  被初始化为0。

8.2 复数 BN PyTorch 实现

定义 self.register_buffer('running_mean', torch.zeros(num_features, dtype = torch.complex64)) 为 BN 的 momentum 的均值；定义 self.register_buffer('running_covar', torch.zeros(num_features,3)) 为 BN 的 momentum 的方差。

class _ComplexBatchNorm(Module):

    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super(_ComplexBatchNorm, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        if self.affine:
            self.weight = Parameter(torch.Tensor(num_features,3))
            self.bias = Parameter(torch.Tensor(num_features,2))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        if self.track_running_stats:
            self.register_buffer('running_mean', torch.zeros(num_features, dtype = torch.complex64))
            self.register_buffer('running_covar', torch.zeros(num_features,3))
            self.running_covar[:,0] = 1.4142135623730951
            self.running_covar[:,1] = 1.4142135623730951
            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        else:
            self.register_parameter('running_mean', None)
            self.register_parameter('running_covar', None)
            self.register_parameter('num_batches_tracked', None)
        self.reset_parameters()

    def reset_running_stats(self):
        if self.track_running_stats:
            self.running_mean.zero_()
            self.running_covar.zero_()
            self.running_covar[:,0] = 1.4142135623730951
            self.running_covar[:,1] = 1.4142135623730951
            self.num_batches_tracked.zero_()

    def reset_parameters(self):
        self.reset_running_stats()
        if self.affine:
            init.constant_(self.weight[:,:2],1.4142135623730951)
            init.zeros_(self.weight[:,2])
            init.zeros_(self.bias)

前向传播时，首先计算当前输入的均值：

mean_r = input.real.mean([0, 2, 3]).type(torch.complex64)
mean_i = input.imag.mean([0, 2, 3]).type(torch.complex64)
mean = mean_r + 1j*mean_i

再进行滑动平均：

if self.training and self.track_running_stats:
# update running mean
with torch.no_grad():
self.running_mean = exponential_average_factor * mean+ (1 - exponential_average_factor) * self.running_mean

输入减去均值：

input = input - mean[None, :, None, None]

计算协方差矩阵的值并做滑动平均：

if self.training or (not self.training and not self.track_running_stats):
# Elements of the covariance matrix (biased for train)
n = input.numel() / input.size(1)
Crr = 1./n*input.real.pow(2).sum(dim=[0,2,3])+self.eps
Cii = 1./n*input.imag.pow(2).sum(dim=[0,2,3])+self.eps
Cri = (input.real.mul(input.imag)).mean(dim=[0,2,3])
else:
Crr = self.running_covar[:,0]+self.eps
Cii = self.running_covar[:,1]+self.eps
Cri = self.running_covar[:,2]#+self.eps

if self.training and self.track_running_stats:
with torch.no_grad():
self.running_covar[:,0] = exponential_average_factor * Crr * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,0]

self.running_covar[:,1] = exponential_average_factor * Cii * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,1]

self.running_covar[:,2] = exponential_average_factor * Cri * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,2]

计算协方差矩阵的逆平方根  ：

det = Crr*Cii-Cri.pow(2)
s = torch.sqrt(det)
t = torch.sqrt(Cii+Crr + 2 * s)
inverse_st = 1.0 / (s * t)
Rrr = (Cii + s) * inverse_st
Rii = (Crr + s) * inverse_st
Rri = -Cri * inverse_st


乘以  再加上  ：

        if self.affine:
            input = (self.weight[None,:,0,None,None]*input.real+self.weight[None,:,2,None,None]*input.imag+\
                    self.bias[None,:,0,None,None]).type(torch.complex64) \
                    +1j*(self.weight[None,:,2,None,None]*input.real+self.weight[None,:,1,None,None]*input.imag+\
                    self.bias[None,:,1,None,None]).type(torch.complex64)

完整的 Complex BN 代码：

class ComplexBatchNorm2d(_ComplexBatchNorm):

def forward(self, input):
exponential_average_factor = 0.0

if self.training and self.track_running_stats:
if self.num_batches_tracked is not None:
self.num_batches_tracked += 1
if self.momentum is None: # use cumulative moving average
exponential_average_factor = 1.0 / float(self.num_batches_tracked)
else: # use exponential moving average
exponential_average_factor = self.momentum

if self.training or (not self.training and not self.track_running_stats):
# calculate mean of real and imaginary part
# mean does not support automatic differentiation for outputs with complex dtype.
mean_r = input.real.mean([0, 2, 3]).type(torch.complex64)
mean_i = input.imag.mean([0, 2, 3]).type(torch.complex64)
mean = mean_r + 1j*mean_i
else:
mean = self.running_mean

if self.training and self.track_running_stats:
# update running mean
with torch.no_grad():
self.running_mean = exponential_average_factor * mean\
+ (1 - exponential_average_factor) * self.running_mean

input = input - mean[None, :, None, None]

if self.training or (not self.training and not self.track_running_stats):
# Elements of the covariance matrix (biased for train)
n = input.numel() / input.size(1)
Crr = 1./n*input.real.pow(2).sum(dim=[0,2,3])+self.eps
Cii = 1./n*input.imag.pow(2).sum(dim=[0,2,3])+self.eps
Cri = (input.real.mul(input.imag)).mean(dim=[0,2,3])
else:
Crr = self.running_covar[:,0]+self.eps
Cii = self.running_covar[:,1]+self.eps
Cri = self.running_covar[:,2]#+self.eps

if self.training and self.track_running_stats:
with torch.no_grad():
self.running_covar[:,0] = exponential_average_factor * Crr * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,0]

self.running_covar[:,1] = exponential_average_factor * Cii * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,1]

self.running_covar[:,2] = exponential_average_factor * Cri * n / (n - 1)\
+ (1 - exponential_average_factor) * self.running_covar[:,2]

# calculate the inverse square root the covariance matrix
det = Crr*Cii-Cri.pow(2)
s = torch.sqrt(det)
t = torch.sqrt(Cii+Crr + 2 * s)
inverse_st = 1.0 / (s * t)
Rrr = (Cii + s) * inverse_st
Rii = (Crr + s) * inverse_st
Rri = -Cri * inverse_st

input = (Rrr[None,:,None,None]*input.real+Rri[None,:,None,None]*input.imag).type(torch.complex64) \
+ 1j*(Rii[None,:,None,None]*input.imag+Rri[None,:,None,None]*input.real).type(torch.complex64)

if self.affine:
input = (self.weight[None,:,0,None,None]*input.real+self.weight[None,:,2,None,None]*input.imag+\
self.bias[None,:,0,None,None]).type(torch.complex64) \
+1j*(self.weight[None,:,2,None,None]*input.real+self.weight[None,:,1,None,None]*input.imag+\
self.bias[None,:,1,None,None]).type(torch.complex64)

return input


9 完整模型搭建

使用复数卷积，BN，激活函数搭建一个简单的完整模型。使用 MNIST 数据集，用文中提到的方法生成虚部。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Subset
from torchvision import datasets, transforms
from complexPyTorch.complexLayers import ComplexBatchNorm2d, ComplexConv2d, ComplexLinear
from complexPyTorch.complexLayers import ComplexDropout2d, NaiveComplexBatchNorm2d
from complexPyTorch.complexLayers import ComplexBatchNorm1d
from complexPyTorch.complexFunctions import complex_relu, complex_max_pool2d

batch_size = 64
n_train = 1000
n_test = 100
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
train_set = datasets.MNIST('../data', train=True, transform=trans, download=True)
train_set = Subset(train_set, torch.arange(n_train))
test_set = datasets.MNIST('../data', train=False, transform=trans, download=True)
test_set = Subset(test_set, torch.arange(n_test))

train_loader = torch.utils.data.DataLoader(train_set, batch_size= batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size= batch_size, shuffle=True)

class ComplexNet(nn.Module):
    
    def __init__(self):
        super(ComplexNet, self).__init__()
        self.conv1 = ComplexConv2d(1, 10, 5, 1)
        self.bn2d  = ComplexBatchNorm2d(10, track_running_stats = False)
        self.conv2 = ComplexConv2d(10, 20, 5, 1)
        self.fc1 = ComplexLinear(4*4*20, 500)
        self.dropout = ComplexDropout2d(p = 0.3)
        self.bn1d  = ComplexBatchNorm1d(500, track_running_stats = False)
        self.fc2 = ComplexLinear(500, 10)
             
    def forward(self,x):
        x = self.conv1(x)
        x = complex_relu(x)
        x = complex_max_pool2d(x, 2, 2)
        x = self.bn2d(x)
        x = self.conv2(x)
        x = complex_relu(x)
        x = complex_max_pool2d(x, 2, 2)
        x = x.view(-1,4*4*20)
        x = self.fc1(x)
        x = self.dropout(x)
        x = complex_relu(x)
        x = self.bn1d(x)
        x = self.fc2(x)
        x = x.abs()
        x =  F.log_softmax(x, dim=1)
        return x
    
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ComplexNet().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=5e-3, momentum=0.9)

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target =data.to(device).type(torch.complex64), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train\t Epoch: {:3} [{:6}/{:6} ({:3.0f}%)]\tLoss: {:.6f}'.format(
                epoch,
                batch_idx * len(data), 
                len(train_loader.dataset),
                100. * batch_idx / len(train_loader), 
                loss.item())
            )
            
def test(model, device, test_loader, optimizer, epoch):
    model.eval()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device).type(torch.complex64), target.to(device)
        output = model(data)
        loss = F.nll_loss(output, target)
        if batch_idx % 100 == 0:
            print('Test\t Epoch: {:3} [{:6}/{:6} ({:3.0f}%)]\tLoss: {:.6f}'.format(
                epoch,
                batch_idx * len(data), 
                len(test_loader.dataset),
                100. * batch_idx / len(test_loader), 
                loss.item())
            )

# Run training on 4 epochs
for epoch in range(4):
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader, optimizer, epoch)

主要参考文献

1 'Deep Complex Networks'

2 论文作者给出的源码地址，使用Theano后端的Keras实现：' https://github.com/ChihebTrabelsi/deep_complex_networks'

3 'https://github.com/wavefrontshaping/complexPyTorch' 给出了部分操作的Pytorch实现版本。

4 深度学习：深度复数网络(Deep Complex Networks)-从论文到pytorch实现:https://www./daima/485c571e9100403

如果觉得有用，就请分享到朋友圈吧！

△点击卡片关注极市平台，获取最新CV干货