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

作者: A君来了 | 来源:发表于2020-01-04 13:21 被阅读0次

    自从深度学习大热以后各种模型层出不穷,但仔细琢磨你会发现,它们无外乎都是由卷积层、全连接层、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这些深度学习框架中。

    相关文章

      网友评论

        本文标题:深入浅出全连接层(fully connected layer)

        本文链接:https://www.haomeiwen.com/subject/fxftnctx.html