使用Pytorch手动实现单层LSTM

背景

最近在研究meta learning（元学习），正在尝试其中一个非常经典的方法叫做MAML（Model Agnostic Meta Learning)。这个方法比较特殊，需要我们维护两套模型的权重，并且进行求导，而Pytorch中封装好的Module模块不支持这样的做法（我没有找到很好的方法，如果有大神知道还望赐教），因此需要手动去实现各种层。

对于卷积、BatchNorm等层，Pytorch提供了Functional接口，如torch.nn.functional.conv2d等，这些接口支持我们传入自定义的权重，因此手动实现一个卷积层相对简单一些，过程如下：

# 定义权重 weight: out_ch, in_ch, kernel_size, kernel_size
# b: out_ch
w = nn.Parameter(torch.ones(out_ch, in_ch, kernel_size, kernel_size))
b = nn.Parameter(torch.zeors(out_ch))
output = nn.functional.conv2d(input, w, b, stride, padding)

然而在官方文档里，并没有找到LSTM的相应接口，所以就必须得自己实现了，这篇文章主要记录一下这个过程。

单层LSTM实现

首先看一下LSTM的原理示意图，此图引自台大李宏毅教授机器学习课程的PPT，关于LSTM比较细致的理解，可以去B站先看一下李洪义老师的课程或者可以看一下知乎陈诚老师的这篇人人都能看懂的LSTM。

本文主要关注于LSTM的实现，对于其原理（比如几个门各自的作用，为什么网络要这么设计之类的）不会有太多阐述，其实主要是因为这一部分我自己还没太弄明白，但是目前觉得并不影响实现，所以如果有错误的地方还望大家指出。

LSTM简单原理

LSTM其实主要就是由隐藏状态h，记忆细胞c和四个门控状态所组成，其中这四个门控状态里可以细分为三个门控状态和一个候选记忆细胞，但是在实现时简单起见可以把他们看成是一样的中间状态。LSTM的过程其实可以看做通过当前时间步输入$x_t$，上一步的隐藏层$h_{t-1}$来决定四个中间状态，然后通过这四个中间状态联合上一步的细胞状态$c_{t-1}$来决定当前时间步的隐藏层$h_t$细胞状态$c_t$以及输出$y_t$

生成的公式在官方文档中有详细的阐述：

LSTM重要向量形状

看明白网络的计算方法后，我们先来看一下网络各个部分比较重要的形状。LSTM要求输入的形状为$T\ast B\ast C$，其中，$T$代表输入的序列长度，也就是时间步的长度；$B$代表Batch size，输入样本的数量；$C$代表特征的通道数，可以理解为一个时间步上的特征用多少个通道来表示。因此可知，每一个时间步的输入$x_t$的形状为$B\ast C$。同时，隐藏状态$h_t$和候选细胞状态$c_{}$