TimeDistributed Pytorch 实现

keras 中的层封装器TimeDistributed模块。官方文档的说明是:
这个封装器将一个层应用于输入的每个时间片。
输入至少为 3D，且第一个维度应该是时间所表示的维度。
考虑 32 个样本的一个 batch， 其中每个样本是 10 个 16 维向量的序列。 那么这个 batch 的输入尺寸为 (32, 10, 16)， 而 input_shape 不包含样本数量的维度，为 (10, 16)。
你可以使用 TimeDistributed 来将 Dense 层独立地应用到 这 10 个时间步的每一个：

# 作为模型第一层
model = Sequential()
model.add(TimeDistributed(Dense(8), input_shape=(10, 16)))
# 现在 model.output_shape == (None, 10, 8)

TimeDistributed 可以应用于任意层，不仅仅是 Dense， 例如运用于 Conv2D 层：

model = Sequential()
model.add(TimeDistributed(Conv2D(64, (3, 3)),
                          input_shape=(10, 299, 299, 3)))

pytorch 中实现相应的封装器。
类似keras中的TimeDistributed封装器，以三维输入为例，分别为(samples, timesteps, input_size)，合并所有的时间步，然后对每个时间步上都进行相同的处理，参数：
module：对每个时间步的进行的操作
batch_first：第一维是否为batch数量

实现代码如下

class TimeDistributed(nn.Module):
    def __init__(self, module, batch_first=False):
        super(TimeDistributed, self).__init__()
        self.module = module
        self.batch_first = batch_first

    def forward(self, x):
        if len(x.size()) <= 2:
            return self.module(x)

        # reshape input data --> (samples * timesteps, input_size)
        # squash timesteps
        x_reshaped = x.contiguous().view(-1, x.size(-1))
        y = self.module(x_reshaped)

        # We have to reshape Y
        if self.batch_first:
            # (samples, timesteps, output_size)
            y = y.contiguous().view(x.size(0), -1, y.size(-1))
        else:
            # (timesteps, samples, output_size)
            y = y.contiguous().view(-1, x.size(1), y.size(-1))
        return y