使用 TensorFlow 进行内存高效滑动 window 序列学习

Question

我正在训练一个自动编码器网络，用于对大型数据集上的多变量时间序列进行编码。我上传了一个完整的示例 gist.

即使它能正常工作，我也必须在内存效率极低或速度缓慢的解决方案之间做出选择。我想更快地获得我的内存高效解决方案。

在我的内存低效设置中，我准备了 plain_dataset 中的训练集。 IE。通过在整个数据集上实现滑动 windows。（在我的真实训练设置中，它们高度重叠）

相反，我想像 idx_dataset 那样将我的训练集定义为 (dataset_index, row_index, size) 元组的列表。在每个训练步骤之前，解析此引用并编译实际训练示例。

我将时间序列数据集转换为 RaggedTensor rt 以便能够以图形模式访问数据。然后我在 resolve_index_batch 中实现了索引解析作为我的 custom_fit 函数的一部分。

我希望与实际训练步骤相比，编写此训练示例的成本相当低，但使用索引训练集时吞吐量几乎减半。

关于如何使 resolve_index_batch 功能更高效的任何想法？

@tf.function
def resolve_index_batch(idx_batch):
    """
    param idx_batch: (32x3) int32 tensor. Each row containing [time_series_idx, row_idx, window_size] 
    """
    return tf.map_fn(fn=lambda idx: tf.slice(rt[idx[0]], [idx[1], 0], [idx[2], -1]), elems=idx_batch, fn_output_signature=tf.float32)

@tf.function
def train_step(batch):
    if mode == 'idx':
        # apply the conversion from indexes batch to data_series batch
        batch = resolve_index_batch(batch)

    # train on reversed time series
    batch_flip = tf.reverse(batch, [1])
    with tf.GradientTape() as tape:
        m = model(batch, training=True)
        loss = loss_fun(batch_flip, m)
        grad = tape.gradient(loss, model.trainable_weights)

    opt.apply_gradients(zip(grad, model.trainable_weights))
    return loss

Answer 1

为了加快resolve_index_batch，用tf.gather代替tf.map_fn。 tf.map_fn只是一个陷阱，很多人掉进去了，在GPU上表现很差。 tf.gather 是向量化操作的正确方法。

示例代码：

@tf.function
def resolve_index_batch_fast(idx_batch):
    """
    note that window length should be a constant, which is the case in your gist codes WINDOW_LEN = 14
    """
    batch_size = idx_batch.shape[0] #if batch size is constant, this line can be optimized as well
    return tf.gather(tf.gather(rt,idx_batch[:,0]),tf.tile(tf.range(window_length)[None,:],[batch_size,1])+idx_batch[:,1:2],batch_dims=1)

在我的实验中，根据批量大小，它至少比 GPU 中的 resolve_index_batch 快数倍。批量越大，加速越大。在 CPU 中，tf.map_fn 效果相当不错。

为了比较普通方法与索引方法之间的性能，拥有这么小的数据并没有那么大的意义，其中整个 rt 可以放入 GPU 内存。在实际和现实的深度学习问题规模中，整个 rt 应该更大，只能驻留在 CPU 内存甚至 SSD 中。

所以首先要做的是通过这样做确保 rt 在 CPU 内存中：

def series_ragged_tensor():
    with tf.device('/CPU:0'):
        rt = tf.RaggedTensor.from_row_lengths(tf.concat(data_series(), axis=0), [ser.shape[0] for ser in data_series()])
    return rt

其次，在CPU中异步做数据准备。里面 def idx_dataset():

return tf.data.Dataset.from_tensor_slices(arr).batch(32).map(resolve_index_batch).repeat().prefetch(tf.data.AUTOTUNE)

第三，将resolve_index_batch定义和rt = series_ragged_tensor()移到正确的位置，相应地使train_step与索引模式和普通模式相同。

使用 TensorFlow 进行内存高效滑动 window 序列学习

Memory efficient sliding window sequence learning with TensorFlow

python

memory

performance

tensorflow