pytorch实现连续切片，类似于tf.strided_slice()

1.背景：

一直使用torch进行开发，torch的友好为开发调试提供了很多方便，但是torch中有些函数没有内置，造成使用的时候必须自己实现，最近就遇到了一个连续切片的问题，问题可以想象我们往Bert中输入的是一个句子对，但是我们还需要第一个句子的表示，因此必须从句子对的表示中取出第一个或者第二个句子的表示，tensorflow中tf.strided_slice()函数有类似的功能，但是我在torch中没有找到类似的函数，因此自己想办法实现了类似功能（如果有大神知道更好的方案请不吝赐教）。

2.解决方案，下面直接上代码：

import torch


# x 表示句子对表示
x = torch.rand((3, 5))
print(x)

# mask是句子1的掩码
mask = [[1, 1, 1, 0],
        [1, 0, 0, 0],
        [1, 1, 0, 0]]
mask = torch.tensor(mask, dtype=torch.long)

# x1是按句子1最大长度截取的表示，可能包含一部分其他信息
x1 = x[:, : mask.size(1)]
print(x1)

# 使用掩码将非句子1对应的表示都变为0
x2 = x1 * mask.float()
print(x2)

函数的参数如下

def strided_slice(input_,
                  begin,
                  end,
                  strides=None,
                  begin_mask=0,
                  end_mask=0,
                  ellipsis_mask=0,
                  new_axis_mask=0,
                  shrink_axis_mask=0,
                  var=None,
                  name=None):

   

具体的理解如下

    input为输入的源数组（数据）
    begin是指的开始的数组下标
    end 是指结束的数组下标
    strides 是指步长

代码实操

input =[
    [[1,1,1],[2,2,2]],
    [[3,3,3],[4,4,4]],
    [[5,5,5],[6,6,6]]
]
r1 = tf.strided_slice(input,[0],[1],[1])
r2 = tf.strided_slice(input,[0],[2],[1])
r3 = tf.strided_slice(input,[0],[2],[2])
with tf.Session() as sess:
    print(sess.run(r1) ## 结果为[[[1,1,1] [2,2,2]]]
    print(sess.run(r2) ##结果为[[[1,1,1] [2,2,2]][[3,3,3],[4,4,4]]
    print(sess.run(r3)) ##结果为[[[1,1,1],[2,2,2]]]

当begin和end只有1位的时候，表示截取最外面的一层

r4 = tf.strided_slice(input,[0,0],[1,1],[1,1])
## 这个表示的就是从最外层截取0到1 就是 [[1,1,1],[2,2,2]] 然后在从里面截取0到1，所以结果是[1,1,1]
r5 = tf.strided_slice(input,[0,0,0],[1,1,1],[1,1,1])
## 这个如下:[[[1,1,1],[2,2,2]]] ⇒ [[[1,1,1]]]==>[[[1]]]

 

当strides大于 end的时候，默认只去begin。当begin>end的时候，取空值
————————————————
版权声明：本文为CSDN博主「aofengdaxia」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/aofengdaxia/article/details/82863058