深入浅出全连接层(fully connected layer)

自从深度学习大热以后各种模型层出不穷,但仔细琢磨你会发现,它们无外乎都是由卷积层、全连接层、pooling、batch norm、activation这几个基本的元素组合而成的。

全连接层指的是层中的每个节点都会连接它下一层的所有节点,它是模仿人脑神经结构来构建的。脑神经科学家们发现,人的认知能力、记忆力和创造力源于不同神经元之间的连接强弱。因此,早期的神经网络一派的创立有点仿生学的意思,以至于到现在还有学者在研究脑神经科学和AI的结合。

实际上,现在这一轮的人工智能早已不是基于仿生学,而是基于统计学,它要解决的是各种数学问题,像是:

本文也将从计算的角度出发,深入浅出全连接层。点击下方链接查看文中完整源码:https://github.com/alexshuang/deep_nerual_network_from_scratch/blob/master/HeadFirstLinearLayer.ipynb

矩阵相乘(GEMM)

全连接层的前向传播过程就是在做矩阵相乘(不考虑bias),input矩阵 * weight矩阵,即C_{cij} = A_{cik} * B_{ckj}:矩阵C的第c个channel第i行第j列的元素 = 矩阵A的第c个channel的第i* 矩阵B的第c个channel的第j列。这个动态演示页面可以帮你快速建立起对矩阵相乘更直观的认识。

Matrix Multiplication
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实现的矩阵乘法函数,算法的时间复杂度是O(N^4),这就是为什么RNN模型训练起来要比CNN模型慢得多,因为相比CNN,RNN的全连接层数太多了。

我们知道numpy有boardcast机制,在这个例子中,可以通过它来去掉matmul()中的K循环,即:

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()的时间复杂度从O(N^4)降低到O(N^3),计算速度提升了17倍。

并行化

numpy的boardcast之所以能去掉K循环,是因为向量的点乘运算(*)是满足结合律和交换律的,因此可以用SIMD指令来并行化计算,即在a_1*b_1 + a_2*b_2 + ... + a_n*b_n = c中,可以用N个cpu core来分别计算a_1 * b_1a_2 * b_2、... a_n * b_n,再将所有计算结果加总起来,这些运算都是同时(并行)进行的,因此比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循环,算法时间复杂度也从O(N^4)变成了O(N),计算速度提升了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这些深度学习框架中。


欢迎关注和点赞,你的鼓励将是我创作的动力

欢迎转发至朋友圈,公众号转载请后台留言申请授权~

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 159,117评论 4 362
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,328评论 1 293
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 108,839评论 0 243
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 44,007评论 0 206
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,384评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,629评论 1 219
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,880评论 2 313
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,593评论 0 198
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,313评论 1 243
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,575评论 2 246
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,066评论 1 260
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,392评论 2 253
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,052评论 3 236
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,082评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,844评论 0 195
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,662评论 2 274
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,575评论 2 270

推荐阅读更多精彩内容