自从深度学习大热以后各种模型层出不穷,但仔细琢磨你会发现,它们无外乎都是由卷积层、全连接层、pooling、batch norm、activation这几个基本的元素组合而成的。
全连接层指的是层中的每个节点都会连接它下一层的所有节点,它是模仿人脑神经结构来构建的。脑神经科学家们发现,人的认知能力、记忆力和创造力源于不同神经元之间的连接强弱。因此,早期的神经网络一派的创立有点仿生学的意思,以至于到现在还有学者在研究脑神经科学和AI的结合。
实际上,现在这一轮的人工智能早已不是基于仿生学,而是基于统计学,它要解决的是各种数学问题,像是:
- 通过“线性层-非线性层”的嵌套组合来构建可以解决任意(近似)问题的数学模型 -- 你真的明白神经网络?
- 通过求损失函数对变量的导数来优化模型的变量 -- 一文搞懂梯度下降&反向传播
- 通过初始化变量来优化变量的梯度消失问题 -- 一文搞懂深度网络初始化
本文也将从计算的角度出发,深入浅出全连接层。点击下方链接查看文中完整源码:https://github.com/alexshuang/deep_nerual_network_from_scratch/blob/master/HeadFirstLinearLayer.ipynb
矩阵相乘(GEMM)
全连接层的前向传播过程就是在做矩阵相乘(不考虑bias),input矩阵 * weight矩阵,即:矩阵的第个channel第行第列的元素 = 矩阵的第个channel的第行 矩阵的第个channel的第列。这个动态演示页面可以帮你快速建立起对矩阵相乘更直观的认识。
M = 128
N = 128
K = 128
C = 32
def matmul(m_a, m_b, m_c):
for c in range(C):
for m in range(M):
for n in range(N):
val = 0.
for k in range(K):
val += m_a[c,m,k] * m_b[c,k,n]
m_c[c,m,n] = val
%time matmul(matrix_a, matrix_b, matrix_c)
CPU times: user 35.1 s, sys: 9.6 ms, total: 35.1 s
Wall time: 35.1 s
matmul()是用python实现的矩阵乘法函数,算法的时间复杂度是,这就是为什么RNN模型训练起来要比CNN模型慢得多,因为相比CNN,RNN的全连接层数太多了。
我们知道numpy有boardcast机制,在这个例子中,可以通过它来去掉matmul()中的循环,即:
def matmul_boardcast(m_a, m_b, m_c):
for c in range(C):
for m in range(M):
for n in range(N):
m_c[c,m,n] = (m_a[c,m,:] * m_b[c,:,n]).sum()
%time matmul_boardcast(matrix_a, matrix_b, matrix_c)
CPU times: user 2.14 s, sys: 78 ms, total: 2.22 s
Wall time: 2.14 s
由于boardcast的作用,matmul_boardcast()的时间复杂度从降低到,计算速度提升了17倍。
并行化
numpy的boardcast之所以能去掉循环,是因为向量的点乘运算()是满足结合律和交换律的,因此可以用SIMD指令来并行化计算,即在中,可以用N个cpu core来分别计算、、... ,再将所有计算结果加总起来,这些运算都是同时(并行)进行的,因此比for循环要快得多。
相比CPU,GPU的并行化能力更强,因为它有成千上万个core。通过matmul_gpu()可以看到,GPU是如何通过并行化来去掉C、M和N这三个for循环的:
kernel = SourceModule("""
__global__ void matmul(double *mat_a, double *mat_b, double *mat_c, int C, int M, int N, int K)
{
/* 将thread id映射到相应的memory address */
int height = blockIdx.y * blockDim.y + threadIdx.y;
int weight = blockIdx.x * blockDim.x + threadIdx.x;
int channel = blockIdx.z * blockDim.z + threadIdx.z;
int thread_idx = channel * M * N + height * N + weight;
/* 通过并行化去掉了for loop C/M/N,只需要for loop K */
if (channel < C && height < M && weight < N) {
double val = 0;
for (int k = 0; k < K; k++)
val += mat_a[channel * M * N + height * N + k] * mat_b[channel * M * N + k * N + weight];
mat_c[thread_idx] = val;
}
}
""")
def matmul_gpu(m_a, m_b, m_c):
dev_a = gpuarray.to_gpu(m_a.reshape(-1))
dev_b = gpuarray.to_gpu(m_b.reshape(-1))
dev_c = gpuarray.to_gpu(m_c.reshape(-1))
matmul_cuda = kernel.get_function("matmul")
matmul_cuda(dev_a, dev_b, dev_c, np.int32(C), np.int32(M), np.int32(N), np.int32(K), block=(32,32,1), grid=(N//32,M//32,C))
return dev_c.get().reshape(C,M,N)
%time c = matmul_gpu(matrix_a, matrix_b, matrix_c)
CPU times: user 7.18 ms, sys: 2 ms, total: 9.18 ms
Wall time: 10.3 ms
我的显卡是Nvidia的,需要用到CUDA(Nvidia GPU编程的开发框架,如果你安装过Tensorflow你会记得它),它的编程语言是C,因此,matmul(double *mat_a, ...)是用C写的。
用“__global__”关键字修饰的函数在GPU编程中称为kernel。GPU有成千上万个core,每个core并行运行同一个kernel,这些kernel通过各自的thread id来明确自己要处理的data,这样的架构就称为SIMD(Single Instruction Multiple Data)。
这里不会展开介绍CUDA和Pycuda编程,如果有兴趣可以学习CUDA_C_Programming_Guide.pdf和相关课程,你只要重点关注下面这部分代码即可:
/* 通过并行化去掉了for loop C/M/N,只需要for loop K */
if (channel < C && height < M && weight < N) {
double val = 0;
for (int k = 0; k < K; k++)
val += mat_a[channel * M * N + height * N + k] * mat_b[channel * M * N + k * N + weight];
mat_c[thread_idx] = val;
}
可以看到,C、M和N这三个for循环已经消失了,只留下K循环,算法时间复杂度也从变成了,计算速度提升了3500倍!
Pytorch Matmul
matmul_gpu()中kernel的效率还可以进一步优化,而这部分工作已经由Nvidia替我们完成了。Nvidia提供了优化好的线性代数运算库--cuBLAS,Pytorch中的matmul()函数会调用它来进行矩阵乘法计算。
matrix_a = torch.randn(C, M, K)
matrix_b = torch.randn(C, K, N)
%time matrix_c = matrix_a.matmul(matrix_b)
CPU times: user 3.38 ms, sys: 0 ns, total: 3.38 ms
Wall time: 2.62 ms
可以看到,Pytorch的matmul()的效率比matmul_gpu()提升了4倍。之所以有400%的提升,是因为cuBLAS充分利用缓存、共享内存、内存局部性等技术,提升了GPU的内存带宽和指令吞吐量,这部分知识与计算机/GPU体系结构相关,在这里就不过多展开了。
小结
全连接层实质上就是矩阵相乘,由于它在数学上满足交换律和结合律,因此可以用并行化来加速计算,cuBLAS就是Nvidia为深度学习提供的数学(矩阵)加速运算库,它已经集成到Pytorch、Tensorflow这些深度学习框架中。
欢迎关注和点赞,你的鼓励将是我创作的动力
欢迎转发至朋友圈,公众号转载请后台留言申请授权~