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

x33g5p2x  于2022-04-07 转载在 其他  
字(1.5k)|赞(0)|评价(0)|浏览(208)

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

相关文章