tensorflow和keras张量切片（slice）

Notes

想将一个向量 $x={\left(x_1,x_2,x_3,x_4,x_5\right)}^T$ 分割成两部分：$a={\left(x_1,x_2\right)}^T$ 和 $b={\left(x_3,x_4,x_5\right)}^T$，操作大概是：
$$a=\left[\begin{array}{cccccc}1& 0& 0& 0& 0& 0\\ 0& 1& 0& 0& 0& 0\end{array}\right]x$$
和
$$b=\left[\begin{array}{cccccc}0& 0& 0& 1& 0& 0\\ 0& 0& 0& 0& 1& 0\\ 0& 0& 0& 0& 0& 1\end{array}\right]x$$
在 TensorFlow 中，用 tf.slice 实现张量切片，Keras 中自定义 Lambda 层实现。

TensorFlow

tf.slice(input_, begin, size, name=None)

• input_：tf.tensor，被操作的 tensor
• begin：list，各个维度的开始下标
• size：list，各个维度上要截多长

import tensorflow as tf

with tf.Session() as sess:
    a = tf.constant([1, 2, 3, 4, 5])
    b = tf.slice(a, [0], [2])  # 第一个维度从 0 开始，截 2 个
    c = tf.slice(a, [2], [3])  # 第一个维度从 2 开始，截 3 个
    print(a.eval())
    print(b.eval())
    print(c.eval())

• 输出

[1 2 3 4 5]
[1 2]
[3 4 5]

Keras

from keras.layers import Lambda
from keras.models import Sequential
import numpy as np

a = np.array([[1, 2, 3, 4, 5]])
model = Sequential([
    Lambda(lambda a: a[:, :2], input_shape=[5])  # 第二维截前 2 个
])

print(model.predict(a))

• 输出

[[1. 2.]]

References

1. TensorFlow分割：tf.slice函数
2. Are there slice layer and split layer in Keras? #890